NanoDet代码逐行精读与修改（四）动态软标签分配：dynamic soft label assigner
笔者已经为nanodet增加了非常详细的注释，代码请戳此仓库：nanodet_detail_notes: detail every detail about nanodet 。
此仓库会跟着文章推送的节奏持续更新！
目录
4. Dynamic Soft Label Assigner
4.1. 初始化和参数
4.2. 筛除不在ground truth中的priors
4.3. 计算损失
4.4. dynamic k matching
4.5. 获得标签分配结果
TODO：编写一个单次label assign的例子和可视化插图
4. Dynamic Soft Label Assigner
随着目标检测网络的发展，大家发现anchor-free和anchor-based、one-stage和two-stage的界限已经十分模糊，而ATSS的发布也指出是否使用anchor和回归效果的好坏并没有太大差别，最关键的是如何为每个prior（可以看作anchor，或者说参考点、回归起点）分配最合适的标签。关于ATSS更详细的内容请参考笔者的这篇博客：anchor-free 模型概览。
ATSS就是一种动态的标签分配方法，它会根据当前预测的结果选出最优的prior对ground truth进行匹配，而不是像之前一样使用先验的固定规则如iou最大、最接近anchor中点、根据尺寸比例等方法进行匹配。由旷视提出的OTA就是将标签分配视作最优传输问题，将ground truth和background当作provider，anchor当作receiver，很好地解决了标签分配中cost计算的问题。再如DETR中的二分图一对一匹配问题，也是一种动态的标签分配方法，需要在训练过程中实时计算cost（关于DETR的介绍请戳：目标检测终章：Vision Transformer
作者在介绍nanodet-plus的文章中也指出：
既然标签匹配需要依赖预测输出，但预测输出又是依赖标签匹配去训练的，但我的模型一开始是随机初始化的，啥也没有呀？那这不就成了一个鸡生蛋，蛋生鸡的问题了吗？由于小模型的检测头非常轻量，在NanoDet中只使用两个深度可分离卷积模块去同时预测分类和回归，和大模型中对分类和回归分别使用4组256channel的3x3卷积来说简直是天壤之别！让这样的检测头从随机初始化的状态去计算Matching Cost做匹配，这是不是有点太难为它了 。
之前的nanodet使用FCOS的方法进行标签分配，但是显然小模型的检测头在训练初期对于位置特征的提取有些力不从心。因此，为了解决小模型初期无法获取较好的预测的问题，作者借鉴了KD（knowledge distillation，关于知识蒸馏的介绍可以看占位符，讲得非常好）的思想，增加了一个AGM模块（已经在第三部分介绍过）并利用AGM的输出进行动态标签分配。
dsl_assigner.py这个模块位于nanodet/model/head/assigner下。
4.1. 初始化和参数
'''
dynamic soft label assigner,根据pred和GT的IOU进行软标签分配
某个pred与GT的IOU越大,最终分配给它的标签值会越接近一,反之会变小
'''
class DynamicSoftLabelAssigner(BaseAssigner):
"""Computes matching between predictions and ground truth with
dynamic soft label assignment.
Args:
topk (int): Select top-k predictions to calculate dynamic k
best matchs for each gt. Default 13.
iou_factor (float): The scale factor of iou cost. Default 3.0.
"""
def __init__(self, topk=13, iou_factor=3.0):
self.topk = topk
self.iou_factor = iou_factor
def assign(
self,
pred_scores,
priors,
decoded_bboxes,
gt_bboxes,
gt_labels,
):
"""Assign gt to priors with dynamic soft label assignment.
Args:
pred_scores (Tensor): Classification scores of one image,
a 2D-Tensor with shape [num_priors, num_classes]
priors (Tensor): All priors of one image, a 2D-Tensor with shape
[num_priors, 4] in [cx, cy, stride_w, stride_y] format.
decoded_bboxes (Tensor): Predicted bboxes, a 2D-Tensor with shape
[num_priors, 4] in [tl_x, tl_y, br_x, br_y] format.
gt_bboxes (Tensor): Ground truth bboxes of one image, a 2D-Tensor
with shape [num_gts, 4] in [tl_x, tl_y, br_x, br_y] format.
gt_labels (Tensor): Ground truth labels of one image, a Tensor
with shape [num_gts].
Returns:
:obj:`AssignResult`: The assigned result.
"""
INF = 100000000
num_gt = gt_bboxes.size(0)
num_bboxes = decoded_bboxes.size(0)
# assign 0 by default
assigned_gt_inds = decoded_bboxes.new_full((num_bboxes,), 0, dtype=torch.long)这里主要是label assign需要的参数，请特别注意几个tensor的维度和长度，之后非常重要！如果读者对于python的切片和索引操作不熟悉的，可能需要去复习一下，AGM这里为了提高速度，使用了大量的tensor索引tensor的操作。下面让我们开始吧！
4.2. 筛除不在ground truth中的priors
# assign 0 by default
assigned_gt_inds = decoded_bboxes.new_full((num_bboxes,), 0, dtype=torch.long)
# 切片得到prior位置(可以看作anchor point的中心点)
prior_center = priors[:, :2]
# 计算prior center到GT左上角和右下角的距离,从而判断prior是否在GT框内
lt_ = prior_center[:, None] - gt_bboxes[:, :2]
rb_ = gt_bboxes[:, 2:] - prior_center[:, None]
deltas = torch.cat([lt_, rb_], dim=-1)
# is_in_gts通过判断deltas全部大于零筛选处在gt中的prior
# [dxlt,dylt,dxrb,dyrb]四个值都需要大于零,则它们中最小的值也要大于零
# tensor.min会返回一个 namedtuple (values, indices),-1代表最后一个维度
# 其中 values 是给定维度 dim 中输入张量的每一行的最小值,并且索引是找到的每个最小值的索引位置
# 这里判断赋值bool,若prior落在gt中对应的[prior_i,gt_j]会变为true
# [i,j]代表第i个prior是否落在第j个ground truth中
is_in_gts = deltas.min(dim=-1).values > 0
# 这一步生成有效prior的索引,这里请注意之所以用sum是因为一个prior可能落在多个GT中
# 因此上一步生成的is_in_gts确定的是某个prior是否落在每一个GT中,只要落在一个GT范围内,便是有效的
valid_mask = is_in_gts.sum(dim=1) > 0
# 利用得到的mask确定由哪些prior生成的pred_box和它们对应的scores是有效的，注意它们的长度
# 注意valid_decoded_bbox和valid_pred_scores的长度是落在gt中prior的个数
# 稍后在dynamic_k_matching()我们会再提到这一点
valid_decoded_bbox = decoded_bboxes[valid_mask]
valid_pred_scores = pred_scores[valid_mask]
num_valid = valid_decoded_bbox.size(0)
# 出现没有预测框或者训练样本中没有GT的情况
if num_gt == 0 or num_bboxes == 0 or num_valid == 0:
# No ground truth or boxes, return empty assignment
max_overlaps = decoded_bboxes.new_zeros((num_bboxes,))
if num_gt == 0:
# No truth, assign everything to background
# 通过这种方式,可以直接在数据集中放置没有标签的图像作为负样本
assigned_gt_inds[:] = 0
if gt_labels is None:
assigned_labels = None
else:
assigned_labels = decoded_bboxes.new_full(
(num_bboxes,), -1, dtype=torch.long
)
return AssignResult(
num_gt, assigned_gt_inds, max_overlaps, labels=assigned_labels
)FCOS式的网络把feature map上的每个格点当作参考点（回归起点），预测得到的数值是距离该参考点的四个数值（上下左右），其做法是将每个落在GT范围内的prior都当作正样本，这同样是一种先验的固定的规则。显然将那些处于ground truth和background边缘的prior直接作为正样本是不太合适的，这里我们先将在gt范围内的priors筛选出来，稍后根据这些priors输出的预测类别和位置算出cost matrix，进一步确定是否要将其当作正样本（并且即使是作为正样本，也会有soft label的衰减），比起原来的方法会合理不少。
4.3. 计算损失
把落在gt范围内的prior筛选出来之后就可以计算IOU loss、class loss和distance loss了。
最终的cost为C_{total}=C_{cls}+\lambda C_{reg}+C_{dis} ， \lambda 为regression cost的调制系数。
其中 C_{reg}=-log(IOU) , C_{dis}=\alpha^{|x_{pred}-x_{gt}|-\beta}。作者提到C_{dis}可以去掉，如果加上可以在训练前期让AGM收敛得更快。
这部分比较难懂的就是tensor的reshape和索引，相信你配着注释看一定能理解。
# 计算有效bbox和gt的iou损失
pairwise_ious = bbox_overlaps(valid_decoded_bbox, gt_bboxes)
# clamp，加上一个很小的数防止出现NaN
iou_cost = -torch.log(pairwise_ious + 1e-7)
# 根据num_valid的数量(有效bbox)生成对应长度的one-hot label,之后用于计算soft lable
# 每个匹配到gt的prior都会有一个[0,0,...,1,0]的tensor,即label位置的元素为一其余为零
gt_onehot_label = (
F.one_hot(gt_labels.to(torch.int64), pred_scores.shape[-1])
.float()
.unsqueeze(0)
.repeat(num_valid, 1, 1)
)
valid_pred_scores = valid_pred_scores.unsqueeze(1).repeat(1, num_gt, 1)
# IOU*onehot得到软标签,直觉上非常好理解,预测框和gt越接近,说明预测的越好
# 那么,稍后计算交叉熵的时候,标签值也会更大
soft_label = gt_onehot_label * pairwise_ious[..., None]
# 计算缩放权重因子,为软标签减去该prior预测的bbox的score
# 可以想象,差距越大说明当前的预测效果越差,稍后的cost计算就应该给一个更大的惩罚
scale_factor = soft_label - valid_pred_scores
# 计算分类交叉熵损失
cls_cost = F.binary_cross_entropy(
valid_pred_scores, soft_label, reduction="none"
) * scale_factor.abs().pow(2.0)
cls_cost = cls_cost.sum(dim=-1)
# 最后得到的匹配开销矩阵,数值为分类损失和IOU损失,这里利用iou_factor作为调制系数
cost_matrix = cls_cost + iou_cost * self.iou_factor4.4. dynamic k matching
接下来就是根据上一部分得到的cost矩阵，进行动态匹配，决定哪些prior最终会得到正样本的监督训练。
def dynamic_k_matching(self, cost, pairwise_ious, num_gt, valid_mask):
"""Use sum of topk pred iou as dynamic k. Refer from OTA and YOLOX.
Args:
cost (Tensor): Cost matrix.
pairwise_ious (Tensor): Pairwise iou matrix.
num_gt (int): Number of gt.
valid_mask (Tensor): Mask for valid bboxes.
"""
matching_matrix = torch.zeros_like(cost)
# select candidate topk ious for dynamic-k calculation
# pairwise_ious匹配成功的组数可能会小于默认的topk,这里取两者最小值防止越界
candidate_topk = min(self.topk, pairwise_ious.size(0))
# 从IOU矩阵中选出IOU最大的topk个匹配
# topk函数返回一个namedtuple(value,indices),indices没用,python里无用变量一般约定用"_"接收
topk_ious, _ = torch.topk(pairwise_ious, candidate_topk, dim=0)
# calculate dynamic k for each gt
# 用topk个预测IOU值之和作为一个GT要分配给prior的个数
# 这个想法很直观,可以把IOU为1看作一个完整目标,那么这些预测框和GT的IOU总和就是最终分配的个数
# clamp规约,最小值为一,因为不可能一个都不给分配,dynmaic_ks的大小和GT个数相同
dynamic_ks = torch.clamp(topk_ious.sum(0).int(), min=1)
# 对每一个GT,挑选出上面计算出的dymamic_ks个拥有最小cost的预测
for gt_idx in range(num_gt):
_, pos_idx = torch.topk(
cost[:, gt_idx], k=dynamic_ks[gt_idx].item(), largest=False # False则返回最小值
)
matching_matrix[:, gt_idx][pos_idx] = 1.0 # 被匹配的(prior,gt)位置上被置1
del topk_ious, dynamic_ks, pos_idx
# 第二个维度是prior,大于1说明一个prior匹配到了多个GT，这里要选出匹配cost最小的GT
prior_match_gt_mask = matching_matrix.sum(1) > 1
if prior_match_gt_mask.sum() > 0:
cost_min, cost_argmin = torch.min(cost[prior_match_gt_mask, :], dim=1)
# 匹配到多个GT的prior的行全部清零
matching_matrix[prior_match_gt_mask, :] *= 0.0
# 把这些prior和gt有最小cost的位置置1
matching_matrix[prior_match_gt_mask, cost_argmin] = 1.0
# get foreground mask inside box and center prior
# 统计matching_matrix中被分配了标签的priors
fg_mask_inboxes = matching_matrix.sum(1) > 0.0
'''
假设priors长度是n,并且有m个prior落在gt中,那么valid_mask长度也是n,并且有m个true,n-m个false;
在这m个落在gt中的prior里,又有k个被匹配到了,故fg_mask_inboxes的维度是m,其中有k个位置为true;
因此valid_mask中有m-k个位置的true也需要被置为false.
在这里valid_mask[valid_mask]会把原来所有为true的位置索引返回,让他们等于fg_mask_inboxes
维度对照表:
n为priors长度即所有prior个数,m为落在gt中的prior个数,k为匹配到gt的prior的个数,g为gt个数
valid_mask:[n] 其中m个为true
matching_matrix:[gt,m] 每一列至多只有一个为1
fg_mask_inboxes:[m] 其中k个为true
最终得到的valid_mask中只有k个为true剩余为false
'''
# 注意这个索引方式,valid_mask是一个bool类型tensor,以自己为索引会返回所有为True的位置
valid_mask[valid_mask.clone()] = fg_mask_inboxes
# 找到已被分配标签的prior对应的gt index,argmax返回最大值所在的索引,每一个prior只会对应一个GT
matched_gt_inds = matching_matrix[fg_mask_inboxes, :].argmax(1)
# 同上,把它们的IOU提取出来
matched_pred_ious = (matching_matrix * pairwise_ious).sum(1)[fg_mask_inboxes]
return matched_pred_ious, matched_gt_inds需要特别注意的有两个地方，因为一个prior可能会匹配到多个GT，当出现这种情况的时候要选择匹配cost最小的那个gt。
第二处是在最后增加了长注释的这一段，务必要清楚，因为valid_mask的长度和cost matrix的长度是不一样的，cost matrix中代表priors的那一维的长度是落在gt中priors的数量，而valid_mask的长度是priors的总数。这里巧妙的利用了valid_mask[valid_mask.clone()]对那些有效的prior进行索引。
4.5. 获得标签分配结果
这部分代码还是在 assign() 函数里，就是调用完 dynamic_k_matching() ，紧接着 4.3。
# 返回值为分配到标签的prior与它们对应的gt的iou和这些prior匹配到的gt索引
matched_pred_ious, matched_gt_inds = self.dynamic_k_matching(
cost_matrix, pairwise_ious, num_gt, valid_mask
)
# convert to AssignResult format
# 把结果还原为priors的长度
assigned_gt_inds[valid_mask] = matched_gt_inds + 1
assigned_labels = assigned_gt_inds.new_full((num_bboxes,), -1)
assigned_labels[valid_mask] = gt_labels[matched_gt_inds].long()
max_overlaps = assigned_gt_inds.new_full(
(num_bboxes,), -INF, dtype=torch.float32
)
max_overlaps[valid_mask] = matched_pred_ious
return AssignResult(
num_gt, assigned_gt_inds, max_overlaps, labels=assigned_labels
)AssignResult 是分配的结果，也被构造成了一个类方便调用和调试。其成员变量和构造请自行参看源码，存储了此次分配中gt的数量 num_gt、分配了prior的gt的索引 assigned_gt_inds和这些gt与prior的iou max_overlaps ，还有标签 labels。