结构最清晰的Yolov3 head和loss实现完全解析

前沿

众所周知，Yolo v3 是一个非常优秀和主流的目标检测算法，各类复现、解读层出不穷。而且又有v4和v5等版本持续发力，但其基本结构和计算逻辑并无太大的变化。mmdetection是一个非常优秀的目标检测开源训练框架，其复现的Yolo v3算法结构非常清晰，实现的颗粒度更细，模块化做的更好，非常适合理解和学习。本文着眼Yolo v3的设计精髓——head和loss部分，结合代码对其实现进行解析，供大家参考。

整体流程

mmdetection中的head设计融合了网络head输出+loss计算+预测结果解析三个部分。其中方法forward()负责网络的原始输出， 方法loss()负责计算loss，方法get_bboxes()负责从网络原始输出解析预测box， 方法forward_train()组 织整个训练的loss计算。而对于Yolo v3的head，其loss计算可以归纳为以下七个步骤：

1. 网络head输出：利用2层卷积操作输出我们想要尺寸的tensor，也是网络原始输出；
2. anchor生成：利用设置的anchor（利用聚类算法，每个分支有3个共9组尺寸的anchor）生成整个特征图上所有的anchor，方便后续计算。
3. gt box网格的分配：gt box按照中心落入那个网格，那个网格负责的原则提前分配好，方便后续计算。
4. 正负样本分配：将全部anchor根据和gt box的iou以及分配的网络，划分为正、负、忽略样本；
5. 样本采样：为了平衡正负样本，按照一定规则（例如随机采样）选择部分anchor进行后续loss计算，yolov3全部采样；
6. gt box编码：将gt box编码为网络输出的相同形式，方便直接计算loss；
7. loss 计算：计算分类、confidence、矩形框位置和宽高的loss，并加权求和最终输出，供计算梯度和反向传播；

   bbox_head=dict(
   type='YOLOV3Head', # Yolo v3 head 类名
   num_classes=80, # 预测类别
   in_channels=[512, 256, 128], # heads输出tensor第一层卷积的输入channel
   out_channels=[1024, 512, 256], # heads输出tensor第一层卷积的输出channel
   anchor_generator=dict(
       type='YOLOAnchorGenerator', # Yolo的anchor生成器
       # anchors，供9个
       base_sizes=[[(116, 90), (156, 198), (373, 326)], [(30, 61), (62, 45), (59, 119)],
       [(10, 13), (16, 30), (33, 23)]],
       strides=[32, 16, 8]), # 输出特征图的stride
       bbox_coder=dict(type='YOLOBBoxCoder'), # bbox的编码器，负责gt box的编码和pred box的解码
   featmap_strides=[32, 16, 8],  # 输出特征图的stride
   loss_cls=dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0,
   reduction='sum'),# 类别loss
   loss_conf=dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=1.0, reduction='sum'), # confidence loss
   loss_xy=dict(type='CrossEntropyLoss', use_sigmoid=True, loss_weight=2.0, reduction='sum'), # box的位置loss
   loss_wh=dict(type='MSELoss', loss_weight=2.0, reduction='sum'))) # box的宽高loss
   # 正负样本分配类，负责所有anchor的正、负、忽略样本的分配
   train_cfg = dict(
   assigner=dict(type='GridAssigner', pos_iou_thr=0.5, neg_iou_thr=0.5, min_pos_iou=0))

网络head输出

这部分比较简单，就是FPN输出的3个分支，通过两层卷积输出预测head。这里最终输出的形式为：batchSize X (5+类别总数) X 特征图宽X特征图宽。其中的5为预测的xywh和confidence。如下图，因为采用的是coco数据集，所以有80个类别，所以输出的tensor的channel输为255，这里假定batchSize为8，网络的输入为尺寸

[8，3，320，320] 的tensor, 。这部分代码是在foward(self, feats)方法中呈现。

anchor 生成

由于Yolo系列都是采用grid cell的方式划分样本位置，因此anchor只有宽高两个属性。预测box的位置(x，y)是相对于其对应grid cell偏移的，其大小是相对anchor的宽高。这点和Faster rcnn以及SSD等算法不一样。mmdetection为兼容两种做法，将anchor的生成统一到相同的形式上来，既利用AnchorGenerator生成特征图上所有的位置的anchor，这里的anchor是有位置属性的。Yolo 的anchor的类为YoloAnchorGenerator。该类主要完成两个任务：1，anchor的生成；2，gt box在grid cell的分配。

anchor生成

这里anchor的表达形式为左上点和右下点，既[x_0,y_0,x_1,y_1]，核心代码如下：

base_anchor = torch.Tensor([
  x_center - 0.5 * w, y_center - 0.5 * h, x_center + 0.5 * w, y_center + 0.5 * h]

这里的x_center和y_center为base的grid cell的中心点坐标，即原图尺度的左上角第一个格子的中心坐标。例如在尺度为20X20的特征图上，其x和y方向的stride均为320/20=16，因此x_center和y_center为[stride_x/2,stride_y/2]=[8,8]。最终获取的一层输出的base_anchors 尺度为3X4，其中3为anchor个数。然后再通过grid_anchors()方法将base_anchors扩充到整个特征图上，为了后续计算方便，对特征图的宽高wh拉成一个维度。最终得到的anchor_list是长度为8（batchsize）的list，list中每一个元素是长度为3（输出层的个数）的list，内包含3个tensor，尺度分别为300X4（3个anchor X 特征图宽10 X 特征图高10，下同），1200X4，4800X4。

gt box 在grid cell中的分配

正如前文所说，Yolo系列按照grid cell来分配样本。gt box的中心点落入哪一个grid cell，哪一个grid cell负责预测该gt box。通过对gt box的分配，最终获取和anchor_list外两层同样尺度的数据，内部tensor长度为特征图宽X特征图高X anchor数目，值为1代表该物体属于该anchor预测（不是真的需要它来负责，下面还会根据iou再次筛选，可以理解为候选anchor）。代码如下：

feat_h, feat_w = featmap_size
# 获取gt的中心位置
gt_bboxes_cx = ((gt_bboxes[:, 0] + gt_bboxes[:, 2]) * 0.5).to(device)
gt_bboxes_cy = ((gt_bboxes[:, 1] + gt_bboxes[:, 3]) * 0.5).to(device)
# 将gt的中心位置映射到特征图尺寸
gt_bboxes_grid_x = torch.floor(gt_bboxes_cx / stride[0]).long()
gt_bboxes_grid_y = torch.floor(gt_bboxes_cy / stride[1]).long()
# 将w和h方向拉成一个维度
gt_bboxes_grid_idx = gt_bboxes_grid_y * feat_w + gt_bboxes_grid_x
# 记录gt所在的grid的mask，存在gt的位置设置为1
responsible_grid = torch.zeros(
  feat_h * feat_w, dtype=torch.uint8, device=device)
responsible_grid[gt_bboxes_grid_idx] = 1
# 将该mask推广到所有的anchor位置
responsible_grid = responsible_grid[:, None].expand(
  responsible_grid.size(0), num_base_anchors).contiguous().view(-1)
return responsible_grid

正负样本分配

该部分做的是确定正负样本，是在anchor维度上的。也就是确定所有的anchor哪些是正样本，哪些是负样本。划分为正样本的anchor意味着负责gt box的预测，训练的时候就会计算gt box的loss。而负样本表明该anchor没有负责任何物体，当然也需要计算loss，但是只计算confidence loss，因为没有目标，所以无法计算box loss 和类别loss。Yolo还有一个设置就是忽略样本，也就是anchor和gt box有较大的iou，但是又不负责预测它，就忽略掉，不计算任何loss。防止有错误的梯度更新到网络，也是为了提高网络的召回率。这里总结如下：

• 正样本：负责预测gt box的anchor。loss计算box loss(包括中心点+宽高)+confidence loss + 类别loss。
• 负样本：不负责预测gt box的anchor。loss只计算confidence loss。
• 忽略样本：和gt box的iou大于一定阈值，但又不负责该gt box的anchor，一般指中心点grid cell附近的其他grid cell 里的anchor。不计算任何loss。

下面看具体实现。代码是同时确定gt box是分配在哪一层的哪一个或几个anchor上。具体的类为GridAssigner，其中输入参数为：Bboxes：所有的anchor。box_responsible_flags：gt 第一步分配的anchor flags，主要是记录在候选anchor中分配。和gt_bboxes。该类遍历batch，维护一个assigned_gt_inds，类似mask的概念，元素值会被分配为-1：忽略样本，0：负样本，正整数：正样本，同时数字代表负责的gt box的索引。具体步骤如下：

第一步，将所有的assigned_gt_inds设置为-1，默认为忽略样本。

第二步，将所有iou小于一定值例如0.5（或者在一定区间的），设置为0，置为负样本。gt box和全部anchor计算iou，这里的boxes为anchor，是带有位置信息的。获取的overlaps 尺度为gt box个数*全部anchor个数（这里为300+1200+4800=6300）。

overlaps = self.iou_calculator(gt_bboxes, bboxes) # 获取全部iou，size为gt个数X6300
max_overlaps, argmax_overlaps = overlaps.max(dim=0) # 找和所有gtbox最大的iou，size为6300，也就是看看每一个anchor，和所有gt box最大的iou有无大过阈值
assigned_gt_inds[(max_overlaps >= 0) & (max_overlaps <= self.neg_iou_thr)] = 0 #如果小于阈值，例如0.5，设置为负样本，不负责任何gt的预测。

第三步，将全部iou中，非负责gt的（记录在box_responsible_flags，非中心点grid cell的anchor）置为-1，该步骤首先排除掉非中心点grid cell的anchor。因为排除掉的部分肯定不是正样本。

#获取和哪一个gt最大的iou，size为6300，和上一步类似，不过获取的都是负责gt box的grid cell里的anchor
max_overlaps, argmax_overlaps = overlaps.max(dim=0)
# 获取的iou和一定阈值对比，例如0.5，大于该值，设置为正样本。
## 可见这一步是将gt box对应的grid cell 里面大于一定阈值的anchor设置为正样本，可能是多个anchor。
pos_inds = (max_overlaps > self.pos_iou_thr) & box_responsible_flags.type(torch.bool)
assigned_gt_inds[pos_inds] = argmax_overlaps[pos_inds] + 1
#------------------------------------------------------------------------------------#
#------------------------------------------------------------------------------------#
#获取全部gt和哪一个anchor最大的iou，尺度为gt的数目，例如有2个gt，那么size就是2
gt_max_overlaps, gt_argmax_overlaps = overlaps.max(dim=1)
# 遍历gt box，找到其最大的anchor，且在负责的grid cell中，设置为正样本。
# 因为上一步，有些gt box并找不到iou大于阈值的anchor，这部分也是要预测的，所以退而求其次，找最大iou的anchor负责它，当然也是在gt box自己的grid cell里的anchor中寻找。
for i in range(num_gts):
  if gt_max_overlaps[i] > self.min_pos_iou:
    if self.gt_max_assign_all:
      max_iou_inds = (overlaps[i, :] == gt_max_overlaps[i]) & \
         box_responsible_flags.type(torch.bool)
      assigned_gt_inds[max_iou_inds] = i + 1

至此，全部anchor全部分配完成，总结一下：

1. 全部anchor，和gt box的iou小于阈值的，设置为负样本；
2. 正样本来自两部分：第一是gt box对应的grid cell里的anchor，iou大于阈值的。第二部分是gt box对应grid cell里的anchor，和gt box iou 最大的那一个；
3. 其余部分，设置为忽略样本；

可以看出，上面2中的第二部分的正样本是最后计算了，因此理论上所有gt box都会分配一个和自己iou最大的anchor。如果预先被2中第一部分分配了，有可能会被其他gt box挤走，也就是标签重写现象。这个以后可以重点分析一下。还可以看出，一个gt box 可以有多个anchor，但是一个anchor只能负责一个gt box。可以理解为，正负样本的分配在训练该样本之前已经做好了，和训练的好坏以及预测的结果并无关系。当然还有另外一种实现方式是：忽略样本由训练过程中的真实预测的box和gt box算iou，较大的且没有被分配到的为忽略样本，是一种动态的分配方式。究竟哪一种方式好我没有去深入思考和测试，知道的小伙伴可以告诉我。

样本采样

在目标检测中，为了保证正负样本平衡，一般采用了采样设置。但通常情况下， Yolov3 所有的样本都有用到，所以采用默认的采样器PseudoSampler，不做任何的采样操作。只是把anchor和gt box 选出来（按照GridAssigner中的信息），这里不再叙述。

gt box编码

分配好正负样本，需要计算loss。因此gt box要和预测的tensor统一到相同的表达上来。经过样本分配和采样操作，最终获取到配对的anchor和gt box，数目是完全相等的。因为为了便于计算，这里将gt box 复制到和样本的anchor相同的数目。如下所示，gt box编码利用self.bbox_coder.encode进行。

# gt box编码
target_map[sampling_result.pos_inds, :4] = self.bbox_coder.encode(
  sampling_result.pos_bboxes, sampling_result.pos_gt_bboxes,
  anchor_strides[sampling_result.pos_inds])
# target的confidence 全部设置为1，v2中采用的是iou，值得注意
target_map[sampling_result.pos_inds, 4] = 1

前面提到过，Yolo根据grid cell分配box的位置，根据anchor的大小预测box宽高，因此mmdetection将gt box或预测box编码和解码的操作抽象出一个类。在yolo_bbox_coder.py中类YOLOBBoxCoder，提供两个方法：encode()和decode()，分别进行gt box的编码和预测box的解码。解码部分是将网络直接预测的值根据anchor还原到gt原图的表达形式，不再叙述。下面是编码方法：

def encode(self, bboxes, gt_bboxes, stride)
    # 作用是将gt box利用grid cell和anchor编码成网络输出的形式，
    # 为了方便和网络的输出直接计算loss。其中bboxes是指anchor，获取gt的中心点和宽高
    x_center_gt = (gt_bboxes[..., 0] + gt_bboxes[..., 2]) * 0.5
    y_center_gt = (gt_bboxes[..., 1] + gt_bboxes[..., 3]) * 0.5
    w_gt = gt_bboxes[..., 2] - gt_bboxes[..., 0]
    h_gt = gt_bboxes[..., 3] - gt_bboxes[..., 1]
    # 获取anchor的中心点和宽高
    x_center = (bboxes[..., 0] + bboxes[..., 2]) * 0.5
    y_center = (bboxes[..., 1] + bboxes[..., 3]) * 0.5
    w = bboxes[..., 2] - bboxes[..., 0]
    h = bboxes[..., 3] - bboxes[..., 1]
    # 计算target
    w_target = torch.log((w_gt / w).clamp(min=self.eps))
    h_target = torch.log((h_gt / h).clamp(min=self.eps))
    # 注意加上0.5的作用是，anchor保存的是相对grid cell中心点box，而网络预测是相对于grid cell的左上角，     # 因此在此上加0.5（做过归一化） 就可以解析到左上角
    x_center_target = ((x_center_gt - x_center) / stride + 0.5).clamp(
      self.eps, 1 - self.eps)
    y_center_target = ((y_center_gt - y_center) / stride + 0.5).clamp(
      self.eps, 1 - self.eps)
    encoded_bboxes = torch.stack(
      [x_center_target, y_center_target, w_target, h_target], dim=-1)
    return encoded_bboxes

loss 计算

至此，所有的anchor全部计算出来并完成了分配，可以直接进行loss的计算了。经过前面的转化，这里遍历所有输出分支（3个）进行loss计算，如下：

# 在样本上计算分类
loss_cls = self.loss_cls(pred_label, target_label, weight=pos_mask)
# 在正+负样本上计算confidence
loss_conf = self.loss_conf(pred_conf, target_conf, weight=pos_and_neg_mask)
# 在正样本上计算中心点损失和宽高损失
loss_xy = self.loss_xy(pred_xy, target_xy, weight=pos_mask)
loss_wh = self.loss_wh(pred_wh, target_wh, weight=pos_mask)

最后将全部loss按照一定的比例加起来构成最终的损失，可以愉快地进行求梯度和反向传播了。

预测流程

最后再补充一下网络预测流程吧。利用测试分支get_bboxes接口，实现逻辑就比较简单了，具体步骤如下：

1. 遍历全部batch中的输出tensor；
2. 利用sigmoid操作将位置x，y预测拉到0到1之间，并利用decode操作获取预测的box；
3. 利用sigmoid操作获取预测的confidence；
4. 利用sigmoid操作获取预测的类别得分；
5. 保留confidence大于一定阈值的部分，对剩下的box进行nms操作，获取最终的box。

结语

本文从head和loss出发，对mmdetection复现的Yolo v3 进行解析。学习Yolo v3实现过程以及模块化代码构造方法。由于本人水平有限，可能理解会有偏差，希望大家指正、赐教。