完整代码百度云直达链接(包含预训练权重)(小白注释)
https://pan.baidu.com/s/1US6e93OaCYOghmF21v0UIA
提取码:z8at
参考链接
【注】代码是大神的代码,在此基础上添加了详细的小白注释,方便我以后阅读。
上一节:详解主干网络darknet53代码
本节代码所在文件
pytorch_yolo3/nets/yolo3.py
pytorch_yolo3/utils/config.py
文章目录
回顾
还记得吗?DarkNet返回三路的feature map,这节就是讲解这三个支路的后续操作,完成了这些操作,yolo3全网也就搭建完毕了。
# 输出三路分支
out3 = self.layer3(x)
out4 = self.layer4(out3)
out5 = self.layer5(out4)
return out3, out4, out5
尺寸分别是out5=(batch_size,1024,13,13)
out4=(batch_size,512,26,26)
out3=(batch_size,256,52,52)
【注】batch_size是你每次放进网络中图片的数量
有了上一节基础,相信大家最pytorch搭建神经网络的基本模块有了一定认识,简单的部分我就略讲了。
搭建yolo_head
这节就是搭建上图绿框部分yolo_head1,yolo_head2,yolo_head3等网络。
温馨提示:)搭建darknet53时顺序是out3、out4、out5,而搭建yolo_head的顺序要从out5的分支yolo_head1开始。
为何要从yolo_head1开始呢,因为它只要一个输入out5,而yolo_head2除了out4输入,还有yolo_head1的第4层输入,yolo_head3类似。先看看yolo_head1、yolo_head2的结构
yolo_head的结构也很工整:1x1的卷积和3x3的卷积交替使用,1x1用来调整通道数(减少数据量),3x3卷积提取特征并增加通道数。有趣的是,整体通道数并没有增加,并且feature map的宽高没有变(3x3卷积步长=1,padding=1),但是这一套操作使得feature map的感受野逐层增大。
搭建yolo_head之前可以把卷积操作打包一下,还记得吗,yolo3里面日常的卷积操作是“卷积+bn+激活”
import torch
import torch.nn as nn
from collections import OrderedDict
from nets.darknet import darknet53
def conv2d(filter_in,filter_out,kernel_size):
pad = (kernel_size-1) // 2 if(kernel_size) else 0
# conv2d 返回一个nn.Sequential,等待一个输入即可调用内部所有模块得到输出
return nn.Sequential(OrderedDict([("conv",nn.Conv2d(filter_in,filter_out,kernel_size=kernel_size,padding=pad,bias=False)),
("bn",nn.BatchNorm2d(filter_out)),
("relu",nn.LeakyReLU(0.1))
]))
完成一个卷积操作,重复使用该操作,只要修改输入、输出通道数就可以搭建yolo_head1了。
def make_yolo_head(in_filters,filters_list,out_filter):
# 把所有层转化成列表方便后面遍历
m = nn.ModuleList([
conv2d(in_filters, filters_list[0], 1),
conv2d(filters_list[0], filters_list[1], 3),
conv2d(filters_list[1], filters_list[0], 1),
conv2d(filters_list[0], filters_list[1], 3),
conv2d(filters_list[1], filters_list[0], 1),
conv2d(filters_list[0], filters_list[1], 3),
# 最后一个卷积就只是一个2D的卷积了(没有bn和激活)
nn.Conv2d(filters_list[1],out_filter,kernel_size=1,stride=1,padding=0,bias=True)
])
搭建yolo3
YoloBody类构造函数需要传入一个config,调用以后再说,这里先看看config的内容:
Config = \
{
"yolo": {
"anchors": [[[116, 90], [156, 198], [373, 326]], # yolo_head1的先验框
[[30, 61], [62, 45], [59, 119]], # yolo_head2的先验框
[[10, 13], [16, 30], [33, 23]]], # yolo_head3的先验框
"classes": 20, # voc数据集的类别数 coco数据集则是80
},
"img_h": 416, # 输入图片的高
"img_w": 416, # 输入图片的宽
}
实际上是一个嵌套的字典config字典里又包含着yolo这么个字典,yolo里面储存着先验框anchors,和类别数量classes。
不了解先验框的话可以参考理解Anchor box的作用
代码如下:
class YoloBody(nn.Moudle):
def __init__(self,config):
super(YoloBody,self).__init__()
self.config = config
# 创建darknet模型,但不导入预训练权重
self.backbone = darknet53(None)
# 终于用上了!!!
# 这三个数是darknet53三条输出支路的输出通道数,即yolo_head的输入通道数
# 1024是yolo_head1的输入通道数;512是yolo_head2的,256是yolo_head3的
# self.layers_out_filters = [256, 512, 1024]
darknet_out_filters = self.backbone.layers_out_filters
# yolo_head的输出通道数 : 3*(5+20) = 75
# 3是先验框的数量,5是x、y、w、h、置信度这5个值
# 20是voc数据集的类别数,80是coco数据集的类别数
# 【注】我贴出来的图是基于coco数据集的,所以通道数:3*(5+80)=255
final_out_filter = len(config["yolo"]["anchors"][0])*(5+config["yolo"]["classes"])
# 搭建yolo_head1、2、3是个nn.ModuleList,待会要遍历用
self.yolo_head1 = make_yolo_head(darknet_out_filters[-1],[512,1024],final_out_filter)
# 搭建yolo_head2,它和yolo_head3都多了卷积+上采样+cat这部分操作
self.yolo_head2_conv = conv2d(512,256,1)
self.yolo_head2_upsample = nn.Upsample(scale_factor=2,mode='nearest')
self.yolo_head2 = make_yolo_head(darknet_out_filters[-2]+256,[256,512],final_out_filter)
# 搭建yolo_head3
self.yolo_head3_conv = conv2d(256, 128, 1)
self.yolo_head3_upsample = nn.Upsample(scale_factor=2, mode='nearest')
self.yolo_head3 = make_yolo_head(darknet_out_filters[-3] + 128, [128, 256], final_out_filter)
def forward(self,x):
# 把yolo_head1第四层卷积结果提出来(要传给yolo_head2)
# 同理,把yolo_head2第四层卷积结果提出来(要传给yolo_head3)
def _branch(yolo_head,head_input):
for i,e in enumerate(yolo_head):
# e就是把yolo_head这个nn.ModuleList里面的各个卷积操作
head_input = e(head_input)
if i==4:
out_branch = head_input
return head_input,out_branch
out3, out4, out5 = self.backbone(x)
# 调用yolo_head1得到输出
final_out1,branch_head2 = _branch(self.yolo_head1,out5)
# 完成卷积+上采样+cat操作后,调用yolo_head2得到输出
head_2in = self.yolo_head2_conv(branch_head2)
head_2in = self.yolo_head2_upsample(head_2in)
# 在第一维cat,也就是通道所在维度,将通道堆叠到一起
head_2in = torch.cat([head_2in,out4],1)
final_out2, branch_head3 = _branch(self.yolo_head2,head_2in)
# 完成卷积+上采样+cat操作后,调用yolo_head3得到输出
head_3in = self.yolo_head3_conv(branch_head3)
head_3in = self.yolo_head3_upsample(head_3in)
# 在第一维cat,也就是通道所在维度,将通道堆叠到一起
head_3in = torch.cat([head_3in, out3], 1)
# 对无用的返回值可以用下划线抑制掉
final_out3, _ = _branch(self.yolo_head3, head_3in)
return final_out1,final_out2,final_out3
# final_out1(batch_size,3*(5+classes),13,13)
# final_out2(batch_size,3*(5+classes),26,26)
# final_out3(batch_size,3*(5+classes),52,52)
如果没说清楚的地方欢迎留言区评论。