YOLOv2 模型原理及代码解析

YOLOv2的改进

1，中心坐标位置预测的改进

YOLOv1 模型预测的边界框中心坐标 $(x,y)$ 是基于 grid 的偏移，这里 grid 的位置是固定划分出来的，偏移量 = 目标位置 - grid 的位置。

边界框的编码过程：YOLOv2 参考了两阶段网络的 anchor boxes 来预测边界框相对先验框的偏移，同时沿用 YOLOv1 的方法预测边界框中心点相对于 grid 左上角位置的相对偏移值。 $(x,y,w,h)$ 的偏移值和实际坐标值的关系如下图所示。

各个字母的含义如下：

• $b_x,b_y,b_w,b_h$ ：模型预测结果转化为 box 中心坐标和宽高后的值
• $t_x,t_y,t_w,t_h$ ：模型要预测的偏移量。
• $c_x,c_y$ ：grid 的左上角坐标，如上图所示。
• $p_w,p_h$ ：anchor 的宽和高，这里的 anchor 是人为定好的一个框，宽和高是固定的。

通过以上定义我们从直接预测位置改为预测一个偏移量，即基于 anchor 框的宽高和 grid 的先验位置的偏移量，位置上使用 grid，宽高上使用 anchor 框，得到最终目标的位置，这种方法叫作 location prediction。

预测偏移不直接预测位置，是因为作者发现直接预测位置会导致神经网络在一开始训练时不稳定，使用偏移量会使得训练过程更加稳定，性能指标提升了 5% 左右。

在数据集的预处理过程中，关键的边界框编码函数如下（代码来自 github，这个版本更清晰易懂）：

def encode(self, boxes, labels, input_size):
    '''Encode target bounding boxes and class labels into YOLOv2 format.
    Args:
        boxes: (tensor) bounding boxes of (xmin,ymin,xmax,ymax) in range [0,1], sized [#obj, 4].
        labels: (tensor) object class labels, sized [#obj,].
        input_size: (int) model input size.
    Returns:
        loc_targets: (tensor) encoded bounding boxes, sized [5,4,fmsize,fmsize].
        cls_targets: (tensor) encoded class labels, sized [5,20,fmsize,fmsize].
        box_targets: (tensor) truth boxes, sized [#obj,4].
    '''
    num_boxes = len(boxes)
    # input_size -> fmsize
    # 320->10, 352->11, 384->12, 416->13, ..., 608->19
    fmsize = (input_size - 320) / 32 + 10
    grid_size = input_size / fmsize

    boxes *= input_size  # scale [0,1] -> [0,input_size]
    bx = (boxes[:,0] + boxes[:,2]) * 0.5 / grid_size  # in [0,fmsize]
    by = (boxes[:,1] + boxes[:,3]) * 0.5 / grid_size  # in [0,fmsize]
    bw = (boxes[:,2] - boxes[:,0]) / grid_size        # in [0,fmsize]
    bh = (boxes[:,3] - boxes[:,1]) / grid_size        # in [0,fmsize]

    tx = bx - bx.floor()
    ty = by - by.floor()

    xy = meshgrid(fmsize, swap_dims=True) + 0.5  # grid center, [fmsize*fmsize,2]
    wh = torch.Tensor(self.anchors)              # [5,2]

    xy = xy.view(fmsize,fmsize,1,2).expand(fmsize,fmsize,5,2)
    wh = wh.view(1,1,5,2).expand(fmsize,fmsize,5,2)
    anchor_boxes = torch.cat([xy-wh/2, xy+wh/2], 3)  # [fmsize,fmsize,5,4]

    ious = box_iou(anchor_boxes.view(-1,4), boxes/grid_size)  # [fmsize*fmsize*5,N]
    ious = ious.view(fmsize,fmsize,5,num_boxes)               # [fmsize,fmsize,5,N]

    loc_targets = torch.zeros(5,4,fmsize,fmsize)  # 5boxes * 4coords
    cls_targets = torch.zeros(5,20,fmsize,fmsize)
    for i in range(num_boxes):
        cx = int(bx[i])
        cy = int(by[i])
        _, max_idx = ious[cy,cx,:,i].max(0)
        j = max_idx[0]
        cls_targets[j,labels[i],cy,cx] = 1

        tw = bw[i] / self.anchors[j][0]
        th = bh[i] / self.anchors[j][1]
        loc_targets[j,:,cy,cx] = torch.Tensor([tx[i], ty[i], tw, th])
    return loc_targets, cls_targets, boxes/grid_size

边界框的解码过程：虽然模型预测的是边界框的偏移量 $(t_x,t_y,t_w,t_h)$，但是可通过以下公式计算出边界框的实际位置。

$$b_x = \sigma(t_x) + c_x \\\\ b_y = \sigma(t_y) + c_y \\\\ b_w = p_{w}e^{t_w} \\\\ b_h = p_{h}e^{t_h}$$

其中， $(c_x, c_y)$ 为 grid 的左上角坐标，因为 $\sigma$ 表示的是 sigmoid 函数，所以边界框的中心坐标会被约束在 grid 内部，防止偏移过多。 $p_w$、 $p_h$ 是先验框（anchors）的宽度与高度，其值相对于特征图大小 $W\times H$ = $13\times 13$ 而言的，因为划分为 $13 \times 13$ 个 grid，所以最后输出的特征图中每个 grid 的长和宽均是 1。知道了特征图的大小，就可以将边界框相对于整个特征图的位置和大小计算出来（均取值 ${0,1}$）。

$$b_x = (\sigma(t_x) + c_x)/W \\\\ b_y = (\sigma(t_y) + c_y)/H \\\\ b_w = p_{w}e^{t_w}/W \\\\ b_h = p_{h}e^{t_h}/H$$

在模型推理的时候，将以上 4 个值分别乘以图片的宽度和长度（像素点值）就可以得到边界框的实际中心坐标和大小。

在模型推理过程中，模型输出张量的解析，即边界框的解码函数如下：

def decode(self, outputs, input_size):
    '''Transform predicted loc/conf back to real bbox locations and class labels.
    Args:
        outputs: (tensor) model outputs, sized [1,125,13,13].
        input_size: (int) model input size.
    Returns:
        boxes: (tensor) bbox locations, sized [#obj, 4].
        labels: (tensor) class labels, sized [#obj,1].
    '''
    fmsize = outputs.size(2)
    outputs = outputs.view(5,25,13,13)

    loc_xy = outputs[:,:2,:,:]   # [5,2,13,13]
    grid_xy = meshgrid(fmsize, swap_dims=True).view(fmsize,fmsize,2).permute(2,0,1)  # [2,13,13]
    box_xy = loc_xy.sigmoid() + grid_xy.expand_as(loc_xy)  # [5,2,13,13]

    loc_wh = outputs[:,2:4,:,:]  # [5,2,13,13]
    anchor_wh = torch.Tensor(self.anchors).view(5,2,1,1).expand_as(loc_wh)  # [5,2,13,13]
    box_wh = anchor_wh * loc_wh.exp()  # [5,2,13,13]

    boxes = torch.cat([box_xy-box_wh/2, box_xy+box_wh/2], 1)  # [5,4,13,13]
    boxes = boxes.permute(0,2,3,1).contiguous().view(-1,4)    # [845,4]

    iou_preds = outputs[:,4,:,:].sigmoid()  # [5,13,13]
    cls_preds = outputs[:,5:,:,:]  # [5,20,13,13]
    cls_preds = cls_preds.permute(0,2,3,1).contiguous().view(-1,20)
    cls_preds = softmax(cls_preds)  # [5*13*13,20]

    score = cls_preds * iou_preds.view(-1).unsqueeze(1).expand_as(cls_preds)  # [5*13*13,20]
    score = score.max(1)[0].view(-1)  # [5*13*13,]
    print(iou_preds.max())
    print(cls_preds.max())
    print(score.max())

    ids = (score>0.5).nonzero().squeeze()
    keep = box_nms(boxes[ids], score[ids])  # NMS 算法去除重复框
    return boxes[ids][keep] / fmsize

2，1 个 gird 只能对应一个目标的改进

或者说很多目标预测不到，查全率低的改进

YOLOv2 首先把 $7 \times 7$ 个区域改为 $13 \times 13$ 个 grid（区域），每个区域有 5 个anchor，且每个 anchor 对应着 1 个类别，那么，输出的尺寸就应该为：[N,13,13,125]

$125 = 5 \times (5 + 20)$

值得注意的是之前 YOLOv1 的每个 grid 只能预测一个目标的分类概率值，两个 boxes 共享这个置信度概率。现在 YOLOv2 使用了 anchor 先验框后，每个 grid 的每个 anchor 都单独预测一个目标的分类概率值。

之所以每个 grid 取 5 个 anchor，是因为作者对 VOC/COCO 数据集进行 K-means 聚类实验，发现当 k=5 时，模型 recall vs. complexity 取得了较好的平衡。当然， $k$ 越好，mAP 肯定越高，但是为了平衡模型复杂度，作者选择了 5 个聚类簇，即划分成 5 类先验框。设置先验框的主要目的是为了使得预测框与 ground truth 的 IOU 更好，所以聚类分析时选用 box 与聚类中心 box 之间的 IOU 值作为距离指标：

$$d(box, centroid) = 1-IOU(box, centroid)$$

与 Faster RCNN 手动设置 anchor 的大小和宽高比不同，YOLOv2 的 anchor 是从数据集中统计得到的。

3，backbone 的改进

作者提出了一个全新的 backbone 网络：Darknet-19，它是基于前人经典工作和该领域常识的基础上进行设计的。Darknet-19 网络和 VGG 网络类似，主要使用 $3 \times 3$ 卷积，并且每个 $2 \times 2$ pooling 操作之后将特征图通道数加倍。借鉴 NIN 网络的工作，作者使用 global average pooling 进行预测，并在 $3 \times 3$ 卷积之间使用 $1 \times 1$ 卷积来降低特征图通道数从而降低模型计算量和参数量。Darknet-19 网络的每个卷积层后面都是用了 BN 层来加快模型收敛，防止模型过拟合。

Darknet-19 网络总共有 19 个卷积层（convolution）、5 最大池化层（maxpooling）。Darknet-19 以 5.58 T的计算量在 ImageNet 数据集上取得了 72.9% 的 top-1 精度和 91.2% 的 top-5 精度。Darket19 网络参数表如下图所示。

检测训练。在 Darknet19 网络基础上进行修改后用于目标检测。首先，移除网络的最后一个卷积层，然后添加滤波器个数为 1024 的 $3 \times 3$ 卷积层，最后添加一个 $1 \times 1$ 卷积层，其滤波器个数为模型检测需要输出的变量个数。对于 VOC 数据集，每个 grid 预测 5 个边界框，每个边界框有 5 个坐标（ $t_x, t_y, t_w, t_h \ 和\ t_o$）和 20 个类别，所以共有 125 个滤波器。我们还添加了从最后的 3×3×512 层到倒数第二层卷积层的直通层，以便模型可以使用细粒度特征。

$$P_r(object)*IOU(b; object) = \sigma (t_o)$$

Yolov2 整个模型结构代码如下：

代码来源 这里。

'''Darknet in PyTorch.'''
import torch
import torch.nn as nn
import torch.nn.init as init
import torch.nn.functional as F

from torch.autograd import Variable


class Darknet(nn.Module):
    # (64,1) means conv kernel size is 1, by default is 3.
    cfg1 = [32, 'M', 64, 'M', 128, (64,1), 128, 'M', 256, (128,1), 256, 'M', 512, (256,1), 512, (256,1), 512]  # conv1 - conv13
    cfg2 = ['M', 1024, (512,1), 1024, (512,1), 1024]  # conv14 - conv18

    def __init__(self):
        super(Darknet, self).__init__()
        self.layer1 = self._make_layers(self.cfg1, in_planes=3)
        self.layer2 = self._make_layers(self.cfg2, in_planes=512)

        #### Add new layers
        self.conv19 = nn.Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1)
        self.bn19 = nn.BatchNorm2d(1024)
        self.conv20 = nn.Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1)
        self.bn20 = nn.BatchNorm2d(1024)
        # Currently I removed the passthrough layer for simplicity
        self.conv21 = nn.Conv2d(1024, 1024, kernel_size=3, stride=1, padding=1)
        self.bn21 = nn.BatchNorm2d(1024)
        # Outputs: 5boxes * (4coordinates + 1confidence + 20classes)
        self.conv22 = nn.Conv2d(1024, 5*(5+20), kernel_size=1, stride=1, padding=0)

    def _make_layers(self, cfg, in_planes):
        layers = []
        for x in cfg:
            if x == 'M':
                layers += [nn.MaxPool2d(kernel_size=2, stride=2, ceil_mode=True)]
            else:
                out_planes = x[0] if isinstance(x, tuple) else x
                ksize = x[1] if isinstance(x, tuple) else 3
                layers += [nn.Conv2d(in_planes, out_planes, kernel_size=ksize, padding=(ksize-1)//2),
                           nn.BatchNorm2d(out_planes),
                           nn.LeakyReLU(0.1, True)]
                in_planes = out_planes
        return nn.Sequential(*layers)

    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = F.leaky_relu(self.bn19(self.conv19(out)), 0.1)
        out = F.leaky_relu(self.bn20(self.conv20(out)), 0.1)
        out = F.leaky_relu(self.bn21(self.conv21(out)), 0.1)
        out = self.conv22(out)
        return out


def test():
    net = Darknet()
    y = net(Variable(torch.randn(1,3,416,416)))
    print(y.size())  # 模型最后输出张量大小 [1,125,13,13]
    
if __name__ == "__main__":
    test()

4，多尺度训练

YOLOv1 输入图像分辨率为 $449 \times 448$，因为使用了 anchor boxes，所以 YOLOv2 将输入分辨率改为 $416 \times 416$。又因为 YOLOv2 模型中只有卷积层和池化层，所以YOLOv2的输入可以不限于 $416 \times 416$ 大小的图片。为了增强模型的鲁棒性，YOLOv2 采用了多尺度输入训练策略，具体来说就是在训练过程中每间隔一定的 iterations 之后改变模型的输入图片大小。由于 YOLOv2 的下采样总步长为 32，所以输入图片大小选择一系列为 32 倍数的值： $\lbrace 320, 352,...,608 \rbrace$ ，因此输入图片分辨率最小为 $320\times 320$，此时对应的特征图大小为 $10\times 10$（不是奇数），而输入图片最大为 $608\times 608$ ，对应的特征图大小为 $19\times 19$ 。在训练过程，每隔 10 个 iterations 随机选择一种输入图片大小，然后需要修最后的检测头以适应维度变化后，就可以重新训练。

采用 Multi-Scale Training 策略，YOLOv2 可以适应不同输入大小的图片，并且预测出很好的结果。在测试时，YOLOv2 可以采用不同大小的图片作为输入，在 VOC 2007 数据集上的测试结果如下图所示。

损失函数

YOLOv2 的损失函数的计算公式归纳如下

第 2,3 行： $t$ 是迭代次数，即前 12800 步我们计算这个损失，后面不计算了。即前 12800 步我们会优化预测的 $(x,y,w,h)$ 与 anchor 的 $(x,y,w,h)$ 的距离 + 预测的 $(x,y,w,h)$ 与 GT 的 $(x,y,w,h)$ 的距离，12800 步之后就只优化预测的 $(x,y,w,h)$与 GT 的 $(x,y,w,h)$ 的距离，原因是这时的预测结果已经较为准确了，anchor已经满足检测系统的需要，而在一开始预测不准的时候，用上 anchor 可以加速训练。

YOLOv2 的损失函数实现代码如下，损失函数计算过程中的模型预测结果的解码函数和前面的解码函数略有不同，其包含关键部分目标 bbox 的解析。

from __future__ import print_function

import torch
import torch.nn as nn
import torch.nn.functional as F

from torch.autograd import Variable
from utils import box_iou, meshgrid

class YOLOLoss(nn.Module):
    def __init__(self):
        super(YOLOLoss, self).__init__()

    def decode_loc(self, loc_preds):
        '''Recover predicted locations back to box coordinates.
        Args:
          loc_preds: (tensor) predicted locations, sized [N,5,4,fmsize,fmsize].
        Returns:
          box_preds: (tensor) recovered boxes, sized [N,5,4,fmsize,fmsize].
        '''
        anchors = [(1.3221,1.73145),(3.19275,4.00944),(5.05587,8.09892),(9.47112,4.84053),(11.2364,10.0071)]
        N, _, _, fmsize, _ = loc_preds.size()
        loc_xy = loc_preds[:,:,:2,:,:]   # [N,5,2,13,13]
        grid_xy = meshgrid(fmsize, swap_dims=True).view(fmsize,fmsize,2).permute(2,0,1)  # [2,13,13]
        grid_xy = Variable(grid_xy.cuda())
        box_xy = loc_xy.sigmoid() + grid_xy.expand_as(loc_xy)  # [N,5,2,13,13]

        loc_wh = loc_preds[:,:,2:4,:,:]  # [N,5,2,13,13]
        anchor_wh = torch.Tensor(anchors).view(1,5,2,1,1).expand_as(loc_wh)  # [N,5,2,13,13]
        anchor_wh = Variable(anchor_wh.cuda())
        box_wh = anchor_wh * loc_wh.exp()  # [N,5,2,13,13]
        box_preds = torch.cat([box_xy-box_wh/2, box_xy+box_wh/2], 2)  # [N,5,4,13,13]
        return box_preds

    def forward(self, preds, loc_targets, cls_targets, box_targets):
        '''
        Args:
          preds: (tensor) model outputs, sized [batch_size,150,fmsize,fmsize].
          loc_targets: (tensor) loc targets, sized [batch_size,5,4,fmsize,fmsize].
          cls_targets: (tensor) conf targets, sized [batch_size,5,20,fmsize,fmsize].
          box_targets: (list) box targets, each sized [#obj,4].
        Returns:
          (tensor) loss = SmoothL1Loss(loc) + SmoothL1Loss(iou) + SmoothL1Loss(cls)
        '''
        batch_size, _, fmsize, _ = preds.size()
        preds = preds.view(batch_size, 5, 4+1+20, fmsize, fmsize)

        ### loc_loss
        xy = preds[:,:,:2,:,:].sigmoid()   # x->sigmoid(x), y->sigmoid(y)
        wh = preds[:,:,2:4,:,:].exp()
        loc_preds = torch.cat([xy,wh], 2)  # [N,5,4,13,13]

        pos = cls_targets.max(2)[0].squeeze() > 0  # [N,5,13,13]
        num_pos = pos.data.long().sum()
        mask = pos.unsqueeze(2).expand_as(loc_preds)  # [N,5,13,13] -> [N,5,1,13,13] -> [N,5,4,13,13]
        loc_loss = F.smooth_l1_loss(loc_preds[mask], loc_targets[mask], size_average=False)

        ### iou_loss
        iou_preds = preds[:,:,4,:,:].sigmoid()  # [N,5,13,13]
        iou_targets = Variable(torch.zeros(iou_preds.size()).cuda()) # [N,5,13,13]
        box_preds = self.decode_loc(preds[:,:,:4,:,:])  # [N,5,4,13,13]
        box_preds = box_preds.permute(0,1,3,4,2).contiguous().view(batch_size,-1,4)  # [N,5*13*13,4]
        for i in range(batch_size):
            box_pred = box_preds[i]  # [5*13*13,4]
            box_target = box_targets[i]  # [#obj, 4]
            iou_target = box_iou(box_pred, box_target)  # [5*13*13, #obj]
            iou_targets[i] = iou_target.max(1)[0].view(5,fmsize,fmsize)  # [5,13,13]

        mask = Variable(torch.ones(iou_preds.size()).cuda()) * 0.1  # [N,5,13,13]
        mask[pos] = 1
        iou_loss = F.smooth_l1_loss(iou_preds*mask, iou_targets*mask, size_average=False)

        ### cls_loss
        cls_preds = preds[:,:,5:,:,:]  # [N,5,20,13,13]
        cls_preds = cls_preds.permute(0,1,3,4,2).contiguous().view(-1,20)  # [N,5,20,13,13] -> [N,5,13,13,20] -> [N*5*13*13,20]
        cls_preds = F.softmax(cls_preds)  # [N*5*13*13,20]
        cls_preds = cls_preds.view(batch_size,5,fmsize,fmsize,20).permute(0,1,4,2,3)  # [N*5*13*13,20] -> [N,5,20,13,13]
        pos = cls_targets > 0
        cls_loss = F.smooth_l1_loss(cls_preds[pos], cls_targets[pos], size_average=False)

        print('%f %f %f' % (loc_loss.data[0]/num_pos, iou_loss.data[0]/num_pos, cls_loss.data[0]/num_pos), end=' ')
        return (loc_loss + iou_loss + cls_loss) / num_pos

YOLOv2 在 VOC2007 数据集上和其他 state-of-the-art 模型的测试结果的比较如下曲线所示。