手撕yolo3系列——详解yolo3整体网络代码（详细注释）

完整代码百度云直达链接（包含预训练权重）(小白注释)
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)

如果没说清楚的地方欢迎留言区评论。

下一节：手撕yolo3系列——详解train训练代码（详细注释）