pytorch實現網路的訓練和推斷——以yolov3為例
一、簡介
Pytorch是目前非常流行的大規模矩陣計算框架，上手簡易，文件詳盡，最新發表的深度學習領域的論文中有多半是以pytorch框架來實現的，足以看出其易用性和流行度。 這篇文章將以yolov3為例，介紹pytorch中如何實現一個網路的訓練和推斷。
二、Pytorch構建深度學習網路
這一部分主要講解一下，在pytorch中構建一個深度學習網路，需要包含哪些部分，各部分都起了什麼作用。不同的框架的實現方式會有許多不同，但基本都包含這些部分。在以下的講解中我隱去了一些具體的實現細節，如果想詳細瞭解，可以前往Pytorch-YOLOv3這個github瞭解，我的講解程式碼也是以它為基礎改編的，兩個版本配合著看能更好地瞭解和上手。
1。datasets
資料集在網路的訓練過程是必須的。通常在訓練指令碼中，會看到類似下面的這樣一行程式碼。
dataloader
=
torch
。
utils
。
data
。
DataLoader
（
Dataset
（
train_path
））
其中的Dataset是自定義的一個類，train_path是訓練資料集的路徑。Dataset通常定義在命名為datasets的檔案內，當然也有以VOCDataset、COCODataset來命名的，其作用都是相同的，定義一個數據集類，以便pytorch呼叫。下面給出一個Dataset的類定義模板，該模板為yolov3的框架所使用。
class
Dataset
（
Dataset
）：
def
__init__
（
self
，
img_dir
，
label_dir
）：
self
。
img_files
=
glob
。
glob
（
os
。
path
。
join
（
img_dir
，
‘*。*’
））
self
。
label_files
=
glob
。
glob
（
os
。
path
。
join
（
label_dir
，
‘*。*’
））
def
__getitem__
（
self
，
index
）：
# === 圖片 ===
# 讀取圖片
img_path
=
self
。
img_files
［
index
%
len
（
self
。
img_files
）］
。
rstrip
（）
img
=
np
。
array
（
Image
。
open
（
img_path
））
img
=
torch
。
from_numpy
（
img
。
transpose
（
2
，
0
，
1
））
。
float
（）
。
div
（
255
）
# 將numpy。array的格式轉為torch。Tensor格式，並轉換通道
# 影象預處理（可選）
# 做一些諸如pad、resize之類的操作
# === 標籤 ===
# 獲取標籤檔案路徑
label_path
=
self
。
label_files
［
index
%
len
（
self
。
img_files
）］
。
rstrip
（）
# 解析標籤檔案（可選）
# 讀取label_path的檔案然後解析，也可直接返回label_path
return
img
，
label_path
def
__len__
（
self
）：
return
len
（
self
。
img_files
）
基本上所有的Dataset類都會包含init、getitem、len這三個函式，在getitem函式中，一般會包含影象預處理和標籤預處理，也有些是把這兩部分放在外部處理，getitem只獲取影象和標籤檔案路徑，值得注意的是，有不少的框架對getitem進行了過載，所以你可能沒有找到getitem函式，但是有其他函式能代替getitem的作用。
2。models
在DL框架中models是一個最為重要的部分，它實現了整個網路的整體結構和具體細節，在一些通用型的大型專案框架內，通常會把這部分拆分成多個modules進行實現，而在一些小專案裡，models也可能僅僅用一個檔案來實現。這裡我還是以yolov3的models來舉例介紹。
# === 讀取cfg配置檔案 ===
def
create_modules
（
cfg
）：
# 根據配置檔案進行解析
return
module_list
# === yolo層定義 ===
class
YOLOLayer
（
nn
。
Module
）：
def
__init__
（
self
，
cfg
）：
super
（
YOLOLayer
，
self
）
。
__init__
（）
def
forward
（
self
，
x
，
targets
=
None
）：
if
targets
is
not
None
：
# === 訓練階段 ===
# 計算loss，根據輸入的x的結果與targets進行計算，最後得到loss
return
x
，
loss
else
：
# === 推斷階段 ===
# 根據輸入的x計算出預測結果
return
x
# === darknet網路結構定義 ===
class
Darknet
（
nn
。
Module
）：
def
__init__
（
self
，
cfg
）：
super
（
Darknet
，
self
）
。
__init__
（）
self
。
module_list
=
create_modules
（
cfg
）
def
forward
（
self
，
x
，
targets
=
None
）：
losses
=
［］
for
module
in
self
。
module_list
：
if
module
is
not
‘YOLO’
：
x
=
module
（
x
）
else
：
# === 訓練階段 ===
if
is_training
：
x
，
loss
=
module
（
x
，
targets
）
losses
。
append
（
loss
）
# === 推斷階段 ===
else
：
x
=
module
（
x
）
return
x
，
losses
在網路結構和yolo層定義中，init和forward這兩個函式是必須的，事實上這兩個函式也是torch內建已經定義過了的，這裡這樣寫實際是過載了這兩個函式。有些專案中可能會把訓練和推斷的forward函式拆分成兩個函式，函式名字也改變了，實際運用時要注意。
3。train
訓練指令碼可以說是網路中最為關鍵的部分，它直接影響了模型的效能和魯棒性。基本上不同網路的訓練指令碼均有不同之處，但是均可以達到一定的效果。一個訓練指令碼一般包含dataloader、optimizer、model三個部分，運用這三個部分構成train迭代迴圈過程。
# 構建model，模型結構
model
=
Darknet
（
model_config_path
）
model
。
apply
（
weights_init_normal
）
model
。
train
（）
if
cuda
：
model
=
model
。
cuda
（）
# 設定dataloader ，資料集載入器
# batch_size根據視訊記憶體大小調整，shuffle是指是否打亂資料集的讀取順序，num_workers是指用多少個執行緒讀取資料集
dataloader
=
torch
。
utils
。
data
。
DataLoader
（
Dataset
（
img_dir
，
label_dir
），
batch_size
=
16
，
shuffle
=
True
，
num_workers
=
4
）
# 設定optimizer，最佳化器
# 最佳化器的種類有非常多，建議新手使用Adam，因為這是一個自適應調整學習率的最佳化器，不需要設定很多引數
# 如果需要精調模型，或者對這方面比較熟練，可以使用SGD+Momentum最佳化器
optimizer
=
torch
。
optim
。
Adam
（
filter
（
lambda
p
：
p
。
requires_grad
，
model
。
parameters
（）））
# 主迴圈train過程
total_epoch
=
10
for
epoch
in
range
（
total_epoch
）：
for
batch_i
，
（
imgs
，
targets
）
in
enumerate
（
dataloader
）：
# 注意，輸入的影象必須進行通道轉換，我這裡忽略了這個步驟，因為我之前已經在Dataset部分實現了
# 這裡的imgs的shape應該為（B， C， H， W），B為batch_size，C為通道，H為高，W為寬
imgs
。
requires_grad
=
True
# imgs的requires_grad屬性必須為True，而targets的requires_grad屬性為False（預設為False）
if
cuda
：
imgs
=
imgs
。
cuda
（）
targets
=
targets
。
cuda
（）
optimizer
。
zero_grad
（）
_
，
loss
=
model
（
imgs
，
targets
）
loss
。
backward
（）
optimizer
。
step
（）
（
‘epoch：’
，
epoch
，
‘batch：’
，
batch_i
，
‘loss：’
，
loss
。
detach
（）
。
cpu
（）
。
numpy
（））
if
epoch
%
1
==
0
：
torch
。
save
（
model
。
state_dict
（），
‘backup。pth’
）
以上就是訓練指令碼中所包含的基本部分，關於loss的計算，有些project把它放在了forward函數里面，也是沒問題的，只要注意進行計算imgs的requires_grad必須為True就可以了。
4。inference
推斷指令碼相對於訓練來說比較簡單，基本上大同小異，只要模型結構沒錯基本上輸出結果都是相同的。
# 載入模型
model
=
Darknet
（
model_config_path
）
params
=
torch
。
load
（
‘backup。pth’
）
model
。
load_state_dict
（
params
）
model
。
eval
（）
# 讀取圖片和推斷
img
=
cv2
。
imread
（
path
）
img
=
torch
。
from_numpy
（
img
。
transpose
（
2
，
0
，
1
））
。
float
（）
。
div
（
255
）
。
unsqueeze
（
0
）
with
torch
。
no_grad
（）：
out
，
_
=
model
（
img
）
# 處理out，例如進行nms和結果顯示，該部分省略
三、總結
以上就是在pytorch中構建模型和訓練、推斷的主要過程，這個部落格主要目的是幫大家理解這個過程，所以對於一些具體實現細節我沒有給出，想詳細瞭解的可以去Pytorch-YOLOv3這個github上進行了解，後續我有時間也會公佈一個我個人的yolov3的pytorch版本。如有疑問也可以在下面評論，我有空會回覆，謝謝。