文章目录

  • 回顾
  • 搭建yolo_head
  • 搭建yolo3


回顾

还记得吗?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是你每次放进网络中图片的数量

yolov2 pytorch实现_d3

有了上一节基础,相信大家最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的结构

yolov2 pytorch实现_神经网络_02


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。

 

代码如下:

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)