本部分是 从0到1 实现YOLO v3 的第二部分,前两部分主要介绍了YOLO的工作原理,包含的模块的介绍以及如何用pytorch搭建完整的YOLOv3网络结构。本部分主要介绍如何完成YOLO的前馈部分。
本文假设读者已经完成了上部分的阅读,以及对pytorch有一定的了解。
首先在工程目录下新建一个darknet.py
文件,接下来使用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)
这个类是nn.Module
的子类,命名为Darknet
,接着进行一些初始化,给类增加一些新的成员:blocks, net_info
和module_list
。
网络的前馈部分都是在foward
的这个函数中完成的,pytorch会自动调用这个函数,首先,foward
用来完成网络从输入到输出的pipline,其次,将输出的featuemap转换为更容易处理的形式。
def forward(self, x, CUDA):
modules = self.blocks[1:]
outputs = {} #We cache the outputs for the route layer
定义的forward函数如上所示,其包括三个参数,self,输入x,和CUDA,CUDA是一个标志位,true表示使用GPU。
这里,我们使用self.blocks
代替self.blocks[1:]
,因为最后一层并不是前馈网络的一部分,因为route和shortcut层需要前面层的feature map,所以,我们将输出换存在outputs这个字典里,字典的key是layer 的名字。
通过create_modules
函数得到了包含YOLO网络各个模块的module_list
,因此,可以通过迭代的方式取出module_list
中的元素构建称为一个完整的YOLO网络。
网络构建很简单,代码如下:
write = 0 #This is explained a bit later
for i, module in enumerate(modules):
module_type = (module["type"])
write 这个参数稍后在解释,在for循环里,迭代取出modules中的元素,module_type
中是每个模块的名字,包括:convolutional,upsample,route,shortcut
等等,对不同的模块需要做不同的处理,处理代码如下:
if module_type == "convolutional" or module_type == "upsample":
x = self.module_list[i](x)
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的最终输出是包含bounding box 属性的卷积特征图,由单元预测的属性bounding box被相互堆叠在一起。 因此,如果必须访问(5,6)处单元格的第二个边界,那么将不得不通过map [5,6,(5 + C):2 *(5 + C)]对它进行索引。 这种形式对输出处理非常不方便,例如通过对象置信度进行阈值处理,向中心添加网格偏移量(offset),应用anchor等。
另一个问题是,由于检测发生在三个尺度上,所以预测图的尺寸将会不同。 尽管三个特征映射的维度不同,但要对它们执行的输出处理操作是相似的。 不得不在单个张量上进行这些操作,而不是三个单独的张量。
代码实现中,我们定义一个函数predict_transform来解决这些问题。
将transform output
函数定义在util.py
文件里面,并在forward
函数中使用它。
首先扩充util.py
的import部分:
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函数采用检测特征映射并将其变成二维张量,其中张量的每一行对应于bounding box的属性,按以下顺序排列。
对于转化来讲,需要做如下事情:
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)
anchor的尺寸和网络block的高度和宽度等属性有关系。 这些属性描述输入图像的尺寸,比检测图更大(以stride为因子)。 因此,我们必须通过检测特征图的stride来划分anchor。
anchors = [(a[0]/stride, a[1]/stride) for a in anchors]
接下来,需要将输出转化为上一节内容讨论过的形式。先对x,y的坐标和目标得分(objectness score)经过一个sigmiod函数,然后将网络的偏置offset:
#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
然后将anchor应用到bounding box的维度中:
#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]))
我们想要在这里做的最后一件事是将检测映射调整为输入图像的大小。 此处的bounding box属性根据特征映射(例如,13 x 13)确定大小。 如果输入图像是416 x 416,我们将这些属性乘以32(总步长,此处是32, 上一文从0到1实现YOLO中有介绍)。
prediction[:,:,:4] *= stride
return prediction
现在我们已经改变了输出张量,现在我们可以将三个不同尺度的检测图连接成一个大张量。 注意,在我们转换之前这是不可能的,因为不能连接具有不同空间维度的特征映射。 但是现在我们的输出张量仅仅作为一个带有boudning box的tabel,完全可以行的方式连接。
我们的一个障碍是我们无法初始化一个空张量,然后将非空(不同形状)张量连接到它。 因此,我们缓存收集器(保持检测的张量)的初始化,直到我们获得第一个检测映射,然后在我们获得后续检测时连接到映射到它。
注意函数forward中的循环之前的write = 0行。 写入标志用于指示我们是否遇到第一次检测。 如果write为0,则表示收集器尚未初始化。 如果它是1,这意味着收集器已经初始化,我们可以将我们的检测图连接到它。
现在,我们已经使用predict_transform函数自己设定了自己,我们编写了用于在前馈函数forward中处理检测特征映射的代码。
在darknet.py
文件的顶部,添加以下import部分。
from util import *
在forward函数中,添加下面的部分,负责处理yolo模块:
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
到此,网络的前馈部分都完成了,为了测试完成的是否正确,可以先用一张图像测试
输入命令,定义一个test
函数,负责读取一张图像输入网络并得到输出:
得到图片:
wget https://github.com/ayooshkathuria/pytorch-yolo-v3/raw/master/dog-cycle-car.png
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)
该张量的形状为1 x 10647 x 85.第一个维度是批量大小,因为我们使用了单个图像,所以它的大小仅为1。 对于批次中的每个图像,我们都有一个10647 x 85的矩阵。 该矩阵中的每一行代表一个boundingbox。 (4个bbox属性,1个对象评分和80个课堂评分)
此时,我们的网络具有随机权重,并且不会产生正确的输出。 我们需要在我们的网络中加载一个权重文件。 我们将为此使用官方的权重文件,链接:https://pjreddie.com/media/files/yolov3.weights。
如果使用的是linux系统,使用如下命令得到权重:
wget https://pjreddie.com/media/files/yolov3.weights
官方权重文件是包含以串行方式存储的权重的二进制文件。
权重只是以浮动形式存储,我们也不知道哪一些权重属于网络的哪一层。要正确加载权重的化, 我们必须了解权重是如何存储的。
首先,权重只属于两种类型的层,即BN层或卷积层。
网络layer的权重与其在配置文件中的顺序完全相同。
当BN层出现在卷积块中时,不存在偏差。 但是,当没有BN layer 时,偏差“权重”必须从文件中读取。
下图总结了权重如何存储权重。
为了正确的加载权重,我们定义一个load_weights
函数,
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位浮点数。我们可以使用numpy加载权重:
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_normalise是否设置为True。 True和False的情况是不一样的。
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)
如果设置的是False,只需加载卷积层的偏置即可。
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_weights函数来加载你的Darknet对象中的权重。
model = Darknet("cfg/yolov3.cfg")
model.load_weights("yolov3.weights")
随着我们的模型的建立和重量的加载,我们终于可以开始检测对象。 在下一部分中,我们将介绍使用对象置信度阈值和非最大抑制来产生我们最终的目标检测。
PyTorch tutorial
Reading binary files with NumPy
nn.Module, nn.Parameter classes
所有的代码可以从这个链接中得到:
https://github.com/ayooshkathuria/YOLO_v3_tutorial_from_scratch