前一節我們實現了YOLO結構中不同類型的層,這一節我們將用Pytorch來實現整個YOLO結構,定義網絡的前向傳播過程,最終能夠實現給定一張圖片獲得檢測輸出。
這個項目使用python 3.5與Pytorch 0.4進行編寫,官方地址。
必備條件:
- 本教程的part1與part2
- Pytorch的基本知識,包括如何使用nn.Module,nn.Sequential,torch.nn.parameter類構建常規的結構
- 使用Pytorch處理圖像
定義網絡:
下面來為我們的檢測器定義網絡。在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繼承自nn.Modulel類。我們初始化類成員變量blocks,net_info和module_list。
實現網絡前向傳播:
網絡的前向傳播是采用復寫nn.Module類的forward方法來實現的。
forward函數有兩個目的,首先是計算輸出,第二個就是將輸出的檢測特征圖轉換成能夠輕松處理的格式(例如對它們進行變換使得跨多個比例的檢測圖能夠連接在一起,因為它們的維度不同,不經轉換無法連接)。
def forward(self, x, CUDA): modules = self.blocks[1:] outputs = {} #We cache the outputs for the route layer
forward接受三個參數self、input x和CUDA,如果為CUDA真,將使用GPU加速前向過程。
我們從self.blocks[1:]開始迭代,這是因為self.blocks的第一個元素是net塊,不是前向傳播的內容。
Route和Shortcut層需要來自之前層的輸出圖,現在我們通過字典outputs緩存每一層的輸出特征圖,字典的鍵就是層的索引,值為特征圖。
與create_modules函數一樣,遍歷包含網絡模塊的module_list。這里需要注意的是,模塊的添加順序與配置文件相同。這表示我們可以簡單地傳過每個模塊來獲得輸出。
write = 0 #This is explained a bit later for i, module in enumerate(modules): module_type = (module["type"])
Convolutional 和 Upsample 層:
如果模塊為卷積模塊或者是上采樣模塊,前傳直接寫成:
if module_type == "convolutional" or module_type == "upsample": x = self.module_list[i](x)
Route / Shortcut 層:
route的代碼分為兩種情況(見part2)。對於連接兩個特征圖這種情況我們使用torch.cat函數,第二個參數為1。這是因為我們希望在深度維上進行拼接(對於Pytorch來說維度為B*C*H*W,所以深度對應着第1維)。
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]] x = torch.cat((map1, map2), 1)
elif module_type == "shortcut": from_ = int(module["from"]) x = outputs[i-1] + outputs[i+from_]
YOLO(檢測層):
YOLO的輸出是一個卷積特征圖,特征圖的深度方向是邊界框的屬性,cell預測的邊界框屬性彼此堆疊。因此,如果你要訪問(5,6)處cell的第二個邊界框,則必須通過map[5,6,(5+C):2*(5+C)] 來找到值。這種形式對於輸出處理非常不方便,如對象置信度閾值化、向中心坐標添加網格偏移量、應用錨點等。
另一個問題是由於我們在三種尺度上做檢測,預測圖的維度將會不同。雖然這三個feature map的維度不同,但是要對它們執行的輸出處理操作是相似的。最好是在一個張量上做這些運算,而不是在三個張量上。為了解決這些問題,我們引入了函數predict_transform。
轉換輸出:
predict_transform函數在util.py文件中,我們將其導入進來在Darknet類的前向傳播中使用。
在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_transform接收五個參數;prediction(我們的輸出),inp_dim(輸入圖片維度),anchors, num_classes 和 CUDA標志位。
def predict_transform(prediction, inp_dim, anchors, num_classes, CUDA = True):
predict_transform函數接收一個檢測的特征映射將其轉換為一個2維的tensor,這里每行的tensor對應到一個邊界框的屬性,以如下形式:
以下的代碼來做上面的轉換。
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) prediction = prediction.view(batch_size, bbox_attrs*num_anchors, grid_size*grid_size) prediction = prediction.transpose(1,2).contiguous() prediction = prediction.view(batch_size, grid_size*grid_size*num_anchors, bbox_attrs)
錨的尺寸與net塊的高度與寬度一致(使用的是[10,13] , [16,30], [33,23],這個的單位其實就是圖片真實像素值),因此我們必須將錨除以檢測特征圖的步長。
anchors = [(a[0]/stride, a[1]/stride) for a in anchors]
現在,我們需要像Part1中討論的那樣轉換我們的輸出。
對x,y坐標與目標置信度進行sigmoid。
#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) 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]))
最后一步,是將檢測圖放大到輸入圖片相同的尺寸。這里的邊界框屬性是根據特征圖調整大小的。如果輸入圖片是416*416,特征圖是13*13,我們就是將這些屬性乘以步長也就是32.
prediction[:,:,:4] *= stride
這就是全部循環體的內容,在函數的結尾返回預測。
return prediction
Detection層回顧:
現在我們已經對輸出向量進行了轉換,能將三個不同尺寸下的檢測圖拼接組成一個大tensor。轉換之前是做不到這點的,因為我們不能將擁有不同空間維的特征圖拼接在一起。但現在我們的輸出向量就像一個表格一樣,以邊界框作為它的行拼接起來就簡單了。
我們在編程上遇到的一個障礙是不能初始化一個空向量,然后將它與一個非空向量拼接在一起。所以我們等到獲得了第一張檢測圖后才初始化collector(保存檢測的tensor),然后獲得后續檢測時才將特征映射拼接在一起。
注意write=0的標志位位於forward函數循環體之前。write標志位指明了我們是否獲得了第一張檢測圖,write=0也就意味着還沒有獲得第一張檢測圖,當它為1的時候代表collector已經被初始化可以直接拼接檢測圖。
現在我們已經實現了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
最后返回檢測結果
return detections
測試前向傳播過程:
這里有個創建虛擬輸入的函數,我們將把這個輸入到我們的網絡。在寫這個函數之前先將下面這張圖片保存到你的工作路徑。如果你是Linux用戶,請輸入:
wget https://github.com/ayooshkathuria/pytorch-yolo-v3/raw/master/dog-cycle-car.png
在darknet.py的前面定義下面這個函數:
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]
這個tensor的維度為1*10647*85。第一維是批大小,因為我們只有一張圖片所以是1。對於每張圖片,我們都有10647*85的輸出表格。表格的每一行代表了一個邊界框(4個bbox屬性,1個目標置信度,80個類得分)。
到目前為止,我們的網絡都是隨機的權重,所以不會產生正確的輸出。下面我們將導入官方的權重文件。
下載預訓練權重:
將預訓練權重下載到你的檢測器文件夾。可以點這里下載。如果你是Linux用戶:
wget https://pjreddie.com/media/files/yolov3.weights
理解權重文件:
官方的權重文件是以串行方式存儲的二進制文件。
我們必須非常准確地加載權重。因為它僅僅作為浮點數存儲,沒有任何東西可以指導我們它屬於哪個層,如果弄錯了不會出現任何提示,比如說你可能把一個batch norm層的權值加載到卷積層了。因此,我們必須理解權重是如何存儲的。
首先,權重只屬於兩種類型的層,要么是batch norm層,要么是卷積層。
這些層的權重完全按照它們在配置文件中出現的順序存儲。比如說一個卷積后面跟着一個Shortcut塊,然后這個Shortcut塊后面跟着另一個卷積塊,權重文件也是先包含前一個卷積塊的權重,然后是后一個卷積塊的權重。
當batch norm層出現在卷積塊中時,沒有偏置參數。但是當沒有batch norm層時,必須從文件中讀取偏置的“權重”。
下圖總結了權重是如何存儲的。
加載權重
讓我們來寫程序加載權重吧!它作為Darknet類的成員函數。除了self之外還接收權重文件的路徑。
def load_weights(self, weightfile):
開頭的160個字節儲存了5個int32的值,組成了文件頭。
#Open the weights file fp = open(weightfile, "rb") #The first 5 values are header information # 1. Major version number # 2. Minor Version Number # 3. Subversion number # 4,5. Images seen by the network (during training) header = np.fromfile(fp, dtype = np.int32, count = 5) self.header = torch.from_numpy(header) self.seen = self.header[3]
現在剩下的就代表權重了。權重是以float32或32-bit foat的形式存儲的,現在使用一個np.ndarray加載剩下的權重。
weights = np.fromfile(fp, dtype = np.float32)
接着迭代權重文件,把它們加載到我們網絡的模塊中。
ptr = 0 for i in range(len(self.module_list)): module_type = self.blocks[i + 1]["type"] #If module_type is convolutional load weights #Otherwise ignore.
在循環體內部,我們先檢查卷積塊是否有batch_normalize。基於這個加載權重。
if module_type == "convolutional": model = self.module_list[i] try: batch_normalize = int(self.blocks[i+1]["batch_normalize"]) except: batch_normalize = 0 conv = model[0]
我們保留一個名為ptr的變量來記錄我們位於權重數組的位置。現在,如果batch_normalize為真,我們以如下方式加載權重。
if (batch_normalize): bn = model[1] #Get the number of weights of Batch Norm Layer num_bn_biases = bn.bias.numel() #Load the weights bn_biases = torch.from_numpy(weights[ptr:ptr + num_bn_biases]) ptr += num_bn_biases bn_weights = torch.from_numpy(weights[ptr: ptr + num_bn_biases]) ptr += num_bn_biases bn_running_mean = torch.from_numpy(weights[ptr: ptr + num_bn_biases]) ptr += num_bn_biases bn_running_var = torch.from_numpy(weights[ptr: ptr + num_bn_biases]) ptr += num_bn_biases #Cast the loaded weights into dims of model weights. bn_biases = bn_biases.view_as(bn.bias.data) bn_weights = bn_weights.view_as(bn.weight.data) bn_running_mean = bn_running_mean.view_as(bn.running_mean) bn_running_var = bn_running_var.view_as(bn.running_var) #Copy the data to model bn.bias.data.copy_(bn_biases) bn.weight.data.copy_(bn_weights) bn.running_mean.copy_(bn_running_mean) bn.running_var.copy_(bn_running_var)
如果batch_normalize非真,直接加載卷積層的偏置
else: #Number of biases num_biases = conv.bias.numel() #Load the weights conv_biases = torch.from_numpy(weights[ptr: ptr + num_biases]) ptr = ptr + num_biases #reshape the loaded weights according to the dims of the model weights conv_biases = conv_biases.view_as(conv.bias.data) #Finally copy the data conv.bias.data.copy_(conv_biases)
最后加載卷積層的權重
#Let us load the weights for the Convolutional layers num_weights = conv.weight.numel() #Do the same as above for weights conv_weights = torch.from_numpy(weights[ptr:ptr+num_weights]) ptr = ptr + num_weights conv_weights = conv_weights.view_as(conv.weight.data) conv.weight.data.copy_(conv_weights)
加載權重的函數就此完成,現在你可以調用Darknet對象里面的load_wights函數
model = Darknet("cfg/yolov3.cfg") model.load_weights("yolov3.weights")
這一部分到此結束,通過構建模型加載權重,我們終於能夠檢測目標了。在下一部分,你會學習如何使用目標置信度與非極大值抑制來產生我們最后的檢測集合。