玩转经典目标检测YOLO系列（二）：详解YOLOV2的复现（上）——整体网络架构与前向推理步骤

经典目标检测YOLO系列(二)YOLOV2的复现(1)总体网络架构及前向推理过程

和之前实现的YOLOv1一样，根据《YOLO目标检测》(ISBN:9787115627094)一书，在不脱离YOLOv2的大部分核心理念的前提下，重构一款较新的YOLOv2检测器，来对YOLOV2有更加深刻的认识。

书中源码连接: RT-ODLab: YOLO Tutorial

对比原始YOLOV2网络，主要改进点如下：

• 添加了后续YOLO中使用的neck，即SPPF模块

• 使用普遍用在RetinaNet、FCOS、YOLOX等通用目标检测网络中的解耦检测头（Decoupled head）

• 修改损失函数，分类分支替换为BCE loss，回归分支替换为GIou loss。

• 由基于边界框的正样本匹配策略，改为基于先验框的正样本匹配策略。

对比之前实现的YOLOV1网络，主要改进点：

• 主干网络由ResNet18改为DarkNet19

• 添加先验框机制

• 正样本匹配策略改为：基于先验框的正样本匹配策略

• YOLOv2代码和之前实现的YOLOv1相比，修改之处不多，建议先看之前实现的YOLOv1的相关文章。

1、YOLOv2网络架构

1.1 DarkNet19主干网络

• 使用原版YOLOv2中提出的DarkNet19作为主干网络(backbone)。
• 不同于分类网络，我们去掉网络中的平均池化层以及分类层。DarkNet19网络的下采样倍数为32，一张图片(416×416×3)经过主干网络，得到13×13×1024的特征图。
• 根据官方的做法，DarkNet19需要现在ImageNet数据集上进行预训练。不过，作者提供了DarkNet19在ImageNet数据集上的预训练权重，因此，我们只需要直接加载即可。
• 这里我们不去实现原版YOLOv2中的passthrough层，仅仅输出一个尺度，即c5层。

# RT-ODLab/models/detectors/yolov2/yolov2_backbone.py

import torch
import torch.nn as nn


model_urls = {
    "darknet19": "https://github.com/yjh0410/image_classification_pytorch/releases/download/weight/darknet19.pth",
}


__all__ = ['DarkNet19']


# --------------------- Basic Module -----------------------
class Conv_BN_LeakyReLU(nn.Module):
    def __init__(self, in_channels, out_channels, ksize, padding=0, stride=1, dilation=1):
        super(Conv_BN_LeakyReLU, self).__init__()
        self.convs = nn.Sequential(
            nn.Conv2d(in_channels, out_channels, ksize, padding=padding, stride=stride, dilation=dilation),
            nn.BatchNorm2d(out_channels),
            nn.LeakyReLU(0.1, inplace=True)
        )

    def forward(self, x):
        return self.convs(x)


# --------------------- DarkNet-19 -----------------------
class DarkNet19(nn.Module):
    def __init__(self):
        
        super(DarkNet19, self).__init__()
        # backbone network : DarkNet-19
        # output : stride = 2, c = 32
        self.conv_1 = nn.Sequential(
            Conv_BN_LeakyReLU(3, 32, 3, 1),
            nn.MaxPool2d((2,2), 2),
        )

        # output : stride = 4, c = 64
        self.conv_2 = nn.Sequential(
            Conv_BN_LeakyReLU(32, 64, 3, 1),
            nn.MaxPool2d((2,2), 2)
        )

        # output : stride = 8, c = 128
        self.conv_3 = nn.Sequential(
            Conv_BN_LeakyReLU(64, 128, 3, 1),
            Conv_BN_LeakyReLU(128, 64, 1),
            Conv_BN_LeakyReLU(64, 128, 3, 1),
            nn.MaxPool2d((2,2), 2)
        )

        # output : stride = 8, c = 256
        self.conv_4 = nn.Sequential(
            Conv_BN_LeakyReLU(128, 256, 3, 1),
            Conv_BN_LeakyReLU(256, 128, 1),
            Conv_BN_LeakyReLU(128, 256, 3, 1),
        )

        # output : stride = 16, c = 512
        self.maxpool_4 = nn.MaxPool2d((2, 2), 2)
        self.conv_5 = nn.Sequential(
            Conv_BN_LeakyReLU(256, 512, 3, 1),
            Conv_BN_LeakyReLU(512, 256, 1),
            Conv_BN_LeakyReLU(256, 512, 3, 1),
            Conv_BN_LeakyReLU(512, 256, 1),
            Conv_BN_LeakyReLU(256, 512, 3, 1),
        )
        
        # output : stride = 32, c = 1024
        self.maxpool_5 = nn.MaxPool2d((2, 2), 2)
        self.conv_6 = nn.Sequential(
            Conv_BN_LeakyReLU(512, 1024, 3, 1),
            Conv_BN_LeakyReLU(1024, 512, 1),
            Conv_BN_LeakyReLU(512, 1024, 3, 1),
            Conv_BN_LeakyReLU(1024, 512, 1),
            Conv_BN_LeakyReLU(512, 1024, 3, 1)
        )


    def forward(self, x):
        c1 = self.conv_1(x)                    # c1
        c2 = self.conv_2(c1)                   # c2
        c3 = self.conv_3(c2)                   # c3
        c3 = self.conv_4(c3)                   # c3
        c4 = self.conv_5(self.maxpool_4(c3))   # c4
        c5 = self.conv_6(self.maxpool_5(c4))   # c5

        return c5


# --------------------- Fsnctions -----------------------
def build_backbone(model_name='darknet19', pretrained=False):
    if model_name == 'darknet19':
        # model
        model = DarkNet19()
        feat_dim = 1024

    # load weight
    if pretrained:
        print('Loading pretrained weight ...')
        url = model_urls['darknet19']
        # checkpoint state dict
        checkpoint_state_dict = torch.hub.load_state_dict_from_url(
            url=url, map_location="cpu", check_hash=True)
        # model state dict
        model_state_dict = model.state_dict()
        # check
        for k in list(checkpoint_state_dict.keys()):
            if k in model_state_dict:
                shape_model = tuple(model_state_dict[k].shape)
                shape_checkpoint = tuple(checkpoint_state_dict[k].shape)
                if shape_model != shape_checkpoint:
                    checkpoint_state_dict.pop(k)
            else:
                checkpoint_state_dict.pop(k)
                print(k)

        model.load_state_dict(checkpoint_state_dict)

    return model, feat_dim


if __name__ == '__main__':
    import time
    model, feat_dim = build_backbone(pretrained=True)
    x = torch.randn(1, 3, 416, 416)
    t0 = time.time()

    for layer in model.children():
        x = layer(x)
        print(layer.__class__.__name__, 'output shape:', x.shape)

    # y = model(x)
    t1 = time.time()
    print('Time: ', t1 - t0)

1.2 添加neck

• 和之前实现的YOLOv1一致，选择YOLOV5版本中所用的SPPF模块。
• 代码在RT-ODLab/models/detectors/yolov2/yolov2_neck.py文件中，不在赘述。

1.3 Detection Head网络

• 和之前实现的YOLOv1一致，即使用解耦检测头（Decoupled head）。
• 代码在RT-ODLab/models/detectors/yolov2/yolov1_head.py文件中，不在赘述。

1.4 预测层

• 如下图，由于预测层多了先验框，因此预测层的输出通道的数量略有变化。

        ## 预测层
        # 与YoloV1相比，YoloV2每个网格会预测5个框(VOC数据集)，因此需×5
        self.obj_pred = nn.Conv2d(head_dim, 1 * self.num_anchors, kernel_size=1)
        self.cls_pred = nn.Conv2d(head_dim, num_classes * self.num_anchors, kernel_size=1)
        self.reg_pred = nn.Conv2d(head_dim, 4 * self.num_anchors, kernel_size=1)

1.5 改进YOLOv2的详细网络图

• 与之前实现的YOLOv1相比，主干网络由ResNet18变为DarkNet19，每个网格预测5个anchor box，其他方面一致。
• 与原版的YOLOv2相比，做了更加符合当下的设计理念的修改，包括添加Neck模块、修改检测头等，但是没有引入passthrough层。
• 尽管和原版的YOLOv2有所差别，但内核思想是一致的，均是在YOLOv1的单级检测架构上引入了先验框。

# RT-ODLab/models/detectors/yolov2/yolov2.py

import torch
import torch.nn as nn
import numpy as np

from utils.misc import multiclass_nms

from .yolov2_backbone import build_backbone
from .yolov2_neck import build_neck
from .yolov2_head import build_head


# YOLOv2
class YOLOv2(nn.Module):
    def __init__(self,
                 cfg,
                 device,
                 num_classes=20,
                 conf_thresh=0.01,
                 nms_thresh=0.5,
                 topk=100,
                 trainable=False,
                 deploy=False,
                 nms_class_agnostic=False):
        super(YOLOv2, self).__init__()
        # ------------------- Basic parameters -------------------
        self.cfg = cfg                                 # 模型配置文件
        self.device = device                           # cuda或者是cpu
        self.num_classes = num_classes                 # 类别的数量
        self.trainable = trainable                     # 训练的标记
        self.conf_thresh = conf_thresh                 # 得分阈值
        self.nms_thresh = nms_thresh                   # NMS阈值
        self.topk = topk                               # topk
        self.stride = 32                               # 网络的最大步长
        self.deploy = deploy
        self.nms_class_agnostic = nms_class_agnostic
        # ------------------- Anchor box -------------------
        self.anchor_size = torch.as_tensor(cfg['anchor_size']).float().view(-1, 2) # [A, 2]
        self.num_anchors = self.anchor_size.shape[0]
        
        # ------------------- Network Structure -------------------
        ## 主干网络
        self.backbone, feat_dim = build_backbone(
            cfg['backbone'], trainable&cfg['pretrained'])

        ## 颈部网络
        self.neck = build_neck(cfg, feat_dim, out_dim=512)
        head_dim = self.neck.out_dim

        ## 检测头
        self.head = build_head(cfg, head_dim, head_dim, num_classes)

        ## 预测层
        # 与YoloV1相比，YoloV2每个网格会预测5个框(VOC数据集)，因此需×5
        self.obj_pred = nn.Conv2d(head_dim, 1 * self.num_anchors, kernel_size=1)
        self.cls_pred = nn.Conv2d(head_dim, num_classes * self.num_anchors, kernel_size=1)
        self.reg_pred = nn.Conv2d(head_dim, 4 * self.num_anchors, kernel_size=1)
    

        if self.trainable:
            self.init_bias()


    def init_bias(self):
        # init bias
        init_prob = 0.01
        bias_value = -torch.log(torch.tensor((1. - init_prob) / init_prob))
        nn.init.constant_(self.obj_pred.bias, bias_value)
        nn.init.constant_(self.cls_pred.bias, bias_value)


    def generate_anchors(self, fmp_size):
        pass
        

    def decode_boxes(self, anchors, reg_pred):
        pass


    def postprocess(self, obj_pred, cls_pred, reg_pred, anchors):
        """
        后处理代码，包括topk操作、阈值筛选和非极大值抑制
        """
        pass


    @torch.no_grad()
    def inference(self, x):
        bs = x.shape[0]
        # 主干网络
        feat = self.backbone(x)

        # 颈部网络
        feat = self.neck(feat)

        # 检测头
        cls_feat, reg_feat = self.head(feat)

        # 预测层
        obj_pred = self.obj_pred(reg_feat)
        cls_pred = self.cls_pred(cls_feat)
        reg_pred = self.reg_pred(reg_feat)
        fmp_size = obj_pred.shape[-2:]

        # anchors: [M, 2]
        anchors = self.generate_anchors(fmp_size)

        # 对 pred 的size做一些view调整，便于后续的处理
        # [B, A*C, H, W] -> [B, H, W, A*C] -> [B, H*W*A, C]
        obj_pred = obj_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, 1)                 # [1, 845=13×13×5, 1]
        cls_pred = cls_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, self.num_classes)
        reg_pred = reg_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, 4)

        # 测试时，默认batch是1，
        # 因此，我们不需要用batch这个维度，用[0]将其取走。
        obj_pred = obj_pred[0]       # [H*W*A, 1]
        cls_pred = cls_pred[0]       # [H*W*A, NC]
        reg_pred = reg_pred[0]       # [H*W*A, 4]

        if self.deploy:
            scores = torch.sqrt(obj_pred.sigmoid() * cls_pred.sigmoid())
            bboxes = self.decode_boxes(anchors, reg_pred)
            # [n_anchors_all, 4 + C]
            outputs = torch.cat([bboxes, scores], dim=-1)

            return outputs
        else:
            # post process
            bboxes, scores, labels = self.postprocess(
                obj_pred, cls_pred, reg_pred, anchors)

            return bboxes, scores, labels


    def forward(self, x):
        if not self.trainable:
            return self.inference(x)
        else:
            bs = x.shape[0]
            # 主干网络
            feat = self.backbone(x)

            # 颈部网络
            feat = self.neck(feat)

            # 检测头
            cls_feat, reg_feat = self.head(feat)

            # 预测层
            obj_pred = self.obj_pred(reg_feat)
            cls_pred = self.cls_pred(cls_feat)
            reg_pred = self.reg_pred(reg_feat)
            fmp_size = obj_pred.shape[-2:]

            # A就是Anchor的数量，VOC数据集上设置为5
            # anchors: [M, 2], M = H*W*A
            anchors = self.generate_anchors(fmp_size)

            # 对 pred 的size做一些view调整，便于后续的处理
            # [B, A*C, H, W] -> [B, H, W, A*C] -> [B, H*W*A, C]
            obj_pred = obj_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, 1)
            cls_pred = cls_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, self.num_classes)
            reg_pred = reg_pred.permute(0, 2, 3, 1).contiguous().view(bs, -1, 4)

            # decode bbox
            box_pred = self.decode_boxes(anchors, reg_pred)

            # 网络输出
            outputs = {"pred_obj": obj_pred,                   # (Tensor) [B, M, 1]
                       "pred_cls": cls_pred,                   # (Tensor) [B, M, C]
                       "pred_box": box_pred,                   # (Tensor) [B, M, 4]
                       "stride": self.stride,                  # (Int)
                       "fmp_size": fmp_size                    # (List) [fmp_h, fmp_w]
                       }           
            return outputs

2、YOLOV2的前向推理

在1.5代码中，还遗留几个问题：

1. 如何从边界框偏移量reg_pred解耦出边界框坐标box_pred？
2. 如何实现后处理操作？
3. 如何计算训练阶段的损失？

2.1 解耦边界框坐标

2.1.1 先验框矩阵的生成

YOLOv2网络配置参数如下，我们从中能看到anchor_size变量。这是基于kmeans聚类，在COCO数据集上聚类出的先验框，由于COCO数据集更大、图片更加丰富，因此我们将这几个先验框用在VOC数据集上。

# RT-ODLab/config/model_config/yolov2_config.py
# YOLOv2 Config

yolov2_cfg = {
    # input
    'trans_type': 'ssd',
    'multi_scale': [0.5, 1.5],
    # model
    'backbone': 'darknet19',
    'pretrained': True,
    'stride': 32,  # P5
    'max_stride': 32,
    # neck
    'neck': 'sppf',
    'expand_ratio': 0.5,
    'pooling_size': 5,
    'neck_act': 'lrelu',
    'neck_norm': 'BN',
    'neck_depthwise': False,
    # head
    'head': 'decoupled_head',
    'head_act': 'lrelu',
    'head_norm': 'BN',
    'num_cls_head': 2,
    'num_reg_head': 2,
    'head_depthwise': False,
    'anchor_size': [[17,  25],