本篇文章是《如何使用PyTorch從零開始實現YOLO(v3)目標檢測算法》的第三部分。這系列論文一共有五篇文章,五篇文章的主要內容在下文中有涉及。如果有問題歡迎和我交流~
如何使用PyTorch從零開始實現YOLO(v3)目標檢測算法?
這是從零開始實現YOLO v3檢測器的教程的第三部分。在上一部分,我們使用YOLO的網絡架構實現了YOLO的五個網絡層,在這一部分,我們打算使用PyTorch實現YOLO的網絡架構:當我們輸入一張圖片的時候,可以產生一個輸出。我們的目標是去設計一個網絡的前向傳播。
1 設計網絡
我之前說過,在PyTorch中,我們使用nn.Module類去構建自定義的網絡架構。讓我們爲我們的目標檢測器定義一個網絡吧。在darknet.py 文件中,我們增加了下面的類。
class Darknet(nn.Module):
def __init__(self, cfgfile):
super(Darknet, self).__init__()
self.blocks = parse_cfg(cfgfile)
self.net_info, self.module_list = create_modules(self.blocks)
這裏,我們定義的類的名字是Darknet.class,而且已經繼承了nn.Module這個類。我們用成員blocks, net_info, 和 module_list初試化這個網絡。
2 實現網絡的前向傳播
通過重寫nn.Module類的forward函數來實現網絡的前向傳播。forward有兩個目的。第一,去計算網絡的輸出;第二,對特徵圖進行處理(比如改變它們,去讓跨尺度的檢測映射圖可以得到連接,否則不可能連接,因爲他們的維度不相同)
def forward(self, x, CUDA):
modules = self.blocks[1:]
outputs = [] #We cache the outputs for the route layer
forward輸入了三個參數,self,輸入x和CUDA,如果CUDA爲真,那麼將會使用CPU去加快前向傳播。這裏,我們迭代了self.bolcks[1:]而不是self.bolcks,因爲self.blocks的第一個元素是net網絡塊,它並不是前行傳播的一部分。因爲路由層和捷徑層需要前面網絡層的輸出映射圖,因此我們將每一層的輸出特徵圖都保存在字典outputs中。字典的鍵值就是這個層的目錄,值就是特徵圖。正如create_modules函數,現在我們遍歷module_list,它包含這個網絡的網絡模型。需要注意的是模型存放的順序和他們在配置文件的順序是一致的,這意味着我們可以通過每一個模塊去運行我們的輸入來得到我們的輸出結果。
write = 0 #This is explained a bit later
for i, module in enumerate(modules):
module_type = (module["type"])
2.1 卷積層和上採樣層
如果這個模型是卷積層或者是上採樣層,這就是前向傳播工作的方式
if module_type == "convolutional" or module_type == "upsample":
x = self.module_list[i](x)
2.2 路由層/捷徑層
如果你看一下路由層的代碼,我們就不得不考慮兩種情況(就像第二部分表述的那樣)。一種情況是我們要去連接兩個特徵圖,我們使用torch.cat函數,它的第二個參數置爲1。這是因爲我們要沿着深度(torch.cat())來連接特徵圖(在PyTorch中,卷積層的輸入和輸出已經變爲B * C * H * W',深度和通道的維度對應)。
elif module_type == "route":
layers = module["layers"]
layers = [int(a) for a in layers]
if (layers[0]) > 0:
layers[0] = layers[0] - i
if len(layers) == 1:
x = outputs[i + (layers[0])]
else:
if (layers[1]) > 0:
layers[1] = layers[1] - i
map1 = outputs[i + layers[0]]
map2 = outputs[i + layers[1]]
# 這邊不是兩個特徵圖進行加法運算,而是along depth進行連接,1是指c的索引
x = torch.cat((map1, map2), 1)
elif module_type == "shortcut":
from_ = int(module["from"])
x = outputs[i-1] + outputs[i+from_]
# 這邊提供另外一個思路,列表可以使用倒序進行索引
# x = outputs[-1]+outputs[from_]
2.3 YOLO層(檢測層)
YOLO層的輸出是一個卷積的特徵圖,它包含沿着特徵圖深度方向的邊界框屬性。由一個單元預測的邊界框屬性一個一個的被堆疊在一起。所以,如果你想得到位於(5,6)的單元的第二個邊界框,你會通過map[5,6,(5+c):2*(5+c)]索引得到它(索引三維的張量太麻煩了)。這樣的形式來處理輸出非常的不方便,比如通過置信度確定閾值,爲中心點添加柵格的偏移,使用錨框等等。
另一個問題是,因爲檢測出現三個尺度,預測特徵圖的尺寸也不相同。儘管三個特徵圖尺寸是不一樣的,但是對他們的輸出處理操作是相似的。在一個單獨的張量中對他們進行運算要好,而不是在三個單獨的的張量中要好。爲了解決這個問題,我們介紹了predict_transform函數。
3 改變輸出
函數predict_transform在util.py文件夾下,當我們在Darknet類的forward函數中使用它的時候,我們需要導入這個函數。
在util.py文件的頂部添加導入
from __future__ import division
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import numpy as np
import cv2
predict_transfor接收五個參數:prediction(我們的輸出),inp_dim(輸入圖片的維度),anchors, num_classes和一個可選的CUDA標誌。
def predict_transform(prediction, inp_dim, anchors, num_classes, CUDA = True):
predict_transform函數輸入一個檢測特徵圖並且把它轉成一個二維張量,張量的每一行以下面的順序對應着邊界框的屬性。
下面的是進行上面轉換的代碼。
batch_size = prediction.size(0)
stride = inp_dim // prediction.size(2)
grid_size = inp_dim // stride
bbox_attrs = 5 + num_classes
num_anchors = len(anchors)
# B*C*H*W --> B*C*(H*W) --> B*(H*W)*((num_classes+5)*3) --> B*(H*W*3)*(num_classes+5)
prediction = prediction.view(batch_size, bbox_attrs*num_anchors, grid_size*grid_size)
# 在PyTorh中,經過transpose之後數據就變得離散了,因此需要經過contiguous將其變成連續的數據
prediction = prediction.transpose(1,2).contiguous()
prediction = prediction.view(batch_size, grid_size*grid_size*num_anchors, bbox_attrs)
錨框的維度對應着net網絡塊的高度和寬度屬性。這些屬性描述了輸入圖片的尺寸,它們比檢測特徵圖的尺寸要大(因爲步長的原因)。因此必須錨框要除以檢測特徵圖的步長。
anchors = [(a[0]/stride, a[1]/stride) for a in anchors]
現在,我們要根據在第一部分討論的方程來改變我們的輸出
#Sigmoid the centre_X, centre_Y. and object confidencce
prediction[:,:,0] = torch.sigmoid(prediction[:,:,0])
prediction[:,:,1] = torch.sigmoid(prediction[:,:,1])
prediction[:,:,4] = torch.sigmoid(prediction[:,:,4])
爲中心座標的預測添加柵格的偏移
#Add the center offsets
grid = np.arange(grid_size)
a,b = np.meshgrid(grid, grid)
x_offset = torch.FloatTensor(a).view(-1,1)
y_offset = torch.FloatTensor(b).view(-1,1)
# 可以通過生成式直接將其轉換成FloatTensor數據
# grid_x = torch.FloatTensor([i for i in range(0,width)])
# grid_y = torch.FloatTensor([i for i in range(0,height)])
if CUDA:
x_offset = x_offset.cuda()
y_offset = y_offset.cuda()
x_y_offset = torch.cat((x_offset, y_offset), 1).repeat(1,num_anchors).view(-1,2).unsqueeze(0)
prediction[:,:,:2] += x_y_offset
將錨點應用到邊界框的尺寸上
#log space transform height and the width
anchors = torch.FloatTensor(anchors)
if CUDA:
anchors = anchors.cuda()
anchors = anchors.repeat(grid_size*grid_size, 1).unsqueeze(0)
prediction[:,:,2:4] = torch.exp(prediction[:,:,2:4])*anchors
將sigmoid激活函數應用到種類的得分上
prediction[:,:,5: 5 + num_classes] = torch.sigmoid((prediction[:,:, 5 : 5 + num_classes]))
這裏,我們想要做的最後一件事就就是將檢測特徵圖的尺寸調整到和輸入圖片的尺寸一樣。這裏,邊界特徵圖的屬性是根據特徵圖的大小確定的的(也就是13*13)。如果輸入圖片的尺寸是416*416,我們要將屬性乘以32,或者是乘以變量stride。
prediction[:,:,:4] *= stride
然後結束了這個循環的主體,在函數的最後返回預測值
return prediction
4 重新審視檢測層
現在我們已經改變數據形狀了,我們現在要把三個尺度的特側特徵圖連接成一個大的張量。注意在轉換之前這是不可能的,因爲不能將不同空間尺寸的特徵圖連接成一個張量。但是現在,我們的輸出張量僅僅就像一個表格一樣,它的行存放邊界框屬性,連接就變得非常有可能了。
在處理過程中一個問題是,在我們得到第一個檢測圖之前,我們不能初始化空的張量,然後和一個非空的(不同形狀的)的張量連接在一起(換句話說,只有當我們運行到檢測層的時候,我們才能夠將輸出進行改變)。我們將收集器(保存檢測器的張量)的初始化放在一邊,直到我們得到了第一個檢測特徵圖,當我們得到後續的檢測器的時候,然後再將特徵圖和它連接在一起。
注意,在forward函數中,write=0這行代碼在循環的外面。這個write標誌用來表示我們是否遇到了第一個檢測器。如果write是0,這表示收集器還沒有初始化。如果是1,這意味着這個收集器已經初始化了,我們可以將檢測映射圖和它進行連接。
現在,我們已經準備好了了predict_transform函數,接下來編寫帶來處理forward函數中的檢測特徵圖。
在你的darknet.py文件的頂部添加下面的引入。
from util import *
然後,再forward函數中
elif module_type == 'yolo':
anchors = self.module_list[i][0].anchors
#Get the input dimensions
inp_dim = int (self.net_info["height"])
#Get the number of classes
num_classes = int (module["classes"])
#Transform
x = x.data
x = predict_transform(x, inp_dim, anchors, num_classes, CUDA)
if not write: #if no collector has been intialised.
detections = x
write = 1
else:
detections = torch.cat((detections, x), 1)
outputs[i] = x
注意:cat是將兩個tensor進行連接,dim=0表示橫向,dim=1表示縱向連接
現在,簡單地返回檢測器
return detections
小結:
(1) 在route層中,數據是逐深度(along depth)進行添加的
(2)爲了方便索引,我們需要改變數據的形狀:由三維數據轉換成二維數據
(3)在創建新的張量的時候:要創建FloatTensor張量,因爲在PyTorch中,只能浮點數纔可以進行求導
(4)三個檢測層輸出的數據也是逐深度進行添加的
5 檢驗前向傳播
這裏是創建一個虛擬(dummy)輸入的函數。我們將會把這個輸入傳遞到我們的網絡。在我們編寫這個函數之前,要先將這個圖片保存在你的工作目錄下。如果你是在Linux環境下,鍵入:
wget https://github.com/ayooshkathuria/pytorch-yolo-v3/raw/master/dog-cycle-car.png
現在,在你的darknet.py文件夾的頂部定義像下面這樣(as follows)這樣函數。
def get_test_input():
img = cv2.imread("dog-cycle-car.png")
img = cv2.resize(img, (416,416)) #Resize to the input dimension
img_ = img[:,:,::-1].transpose((2,0,1)) # BGR -> RGB | H X W C -> C X H X W
img_ = img_[np.newaxis,:,:,:]/255.0 #Add a channel at 0 (for batch) | Normalise
img_ = torch.from_numpy(img_).float() #Convert to float
img_ = Variable(img_) # Convert to Variable
return img_
然後,鍵入下面的代碼:
model = Darknet("cfg/yolov3.cfg")
inp = get_test_input()
pred = model(inp, torch.cuda.is_available())
print (pred)
你將會看到這樣的輸出:
( 0 ,.,.) =
16.0962 17.0541 91.5104 ... 0.4336 0.4692 0.5279
15.1363 15.2568 166.0840 ... 0.5561 0.5414 0.5318
14.4763 18.5405 409.4371 ... 0.5908 0.5353 0.4979
⋱ ...
411.2625 412.0660 9.0127 ... 0.5054 0.4662 0.5043
412.1762 412.4936 16.0449 ... 0.4815 0.4979 0.4582
412.1629 411.4338 34.9027 ... 0.4306 0.5462 0.4138
[torch.FloatTensor of size 1x10647x85]
這個張量的形狀是1*10647*85。第一維是批的大小,這裏是1是因爲我們使用了單獨的一張圖片。在批次中的每一張圖片,我們都有一個10647*85的表格( a 10647*85 table)。這個矩陣(table)每一行代表一個邊界框(4個邊界框屬性,1個目標對象得分和80個類別得分)
這個時候,我們的網絡有一個隨機的權重,不會計算出正確的輸出。在我們的網絡中,我們需要加載一個權重文件。我們將會使用官方的權重文件來達到這個目的(如何加載權重文件的內容我就不贅述了,可以參考原文)。
6 總結
在這一節中,我們主要實現了Darknet的前向傳播和爲模型加載官方數據權重。因爲網絡模型都保存在ModuleList中,因此我們需要再forward函數中手動指定數據在層與層之間流通的順序,因此需要遍歷每一個網絡層指定他們的先後順序。因爲,路由層需要用到前面層的輸出特徵圖,因此我們需要使用一個變量output來保存每一個輸出特徵圖也就是x。一共可能會遇到五類網絡層:
- 卷積層和上採樣層:因爲這兩個層是官方提供的層,因此我們直接將x傳入到他們的實例對象,然後返回計算輸出特徵圖。
- 路由層:路由層只不過是返回到某層,或者是兩個層的疊加。因此只需要返回該層的輸出特徵圖,或者某兩層的輸出特徵圖的疊加(會用到torch.cat([map1,map2],1)方法,爲1的時候,只是通道數相加,爲0的時候是批量數進行相加)
- 捷徑層:捷徑層只是將前面特徵圖與捷徑層前一個特徵圖進行對應元素進行相加,因此直接使用加法就行
- YOLO層:也就是檢測層。我們需要對輸出的結果進行處理,比如將tx,ty,to進行sigmoid運算,需要對tw,th進行指數運算。但是如果直接在三維輸出特徵圖上直接進行,索引相應值比較麻煩,而且不利於處理,因此我們會對輸出特徵圖提前進行處理,將三維特徵圖轉換成二維張量,每一行代表一個邊界框,一共3*W*H行(W,H是指特徵圖的寬和高),一共85列代表表示一個邊界框的85個屬性。因爲我們有三個尺度的輸出,我們需要將三個尺度輸出疊加在一起,因此需要用到一個一個參數detections,保存每一個尺度的輸出特徵圖,每次輸出的時候都使用torch.cat方法,將本次的輸出和上次的輸出疊加在一起。
注意:在pedict_transorm函數中,參數的輸入主要是(1)網絡的預測值(2)錨框(3)輸入圖片的尺寸(4)是否使用GPU,函數的主要工作主要是:
- 對tx,ty,to進行sigmoid操作,並且加上中心偏移
- 對tw,th進行指數運算,並且乘以feature map中的錨框的寬和高
- 對80個類進行sigmoid操作
- 將預測的邊界框的四個值轉換到原圖中
注意:要修改配置文件中net的:height和width
易忽視點:
- 我們要將原圖中的邊界框大小除以stride,轉換到feature map相應的尺寸上,然後最後再轉換到實際圖中
其中比較重要的技巧:
- repeat()函數的使用
- 怎麼將一個三維數據轉換成我們需要的二維數據(這是爲了方便索引)
- 怎麼將tx,ty加上他們各自所在柵格的座標(np.meshgride函數,repeat函數)
- 怎麼將pw,ph乘在相應的元素上面(repeat函數)
- 注意numpy數據類型-->FloatTensor()轉換成tensor --> cuda如果有的話