nanodet-plus阅读：（2）正负样本定义(SimOTA)

一、前言

介绍几篇写的不错的博客，关于SimOTA的。其实它们已经写的很好了，但我还是狗尾续貂般写了这篇博客，这主要是为了方便自己复习，毕竟人总是更加习惯自己写的东西。另外贴上我写的nanodet-plus simOTA代码注释与我认为不错的博客：
1. 动态软标签分配；
2. Yolox核心基础完整讲解；
3. 解读YOLOX（OTA / SimOTA)；

二、正文

首先用语言来描述一下SimOTA的流程，我尽量把我看代码时的疑问解释清楚。

1. 判断每个prior（其实就是YOLO系列的cell/grid，作用类似于anchor）的左上角是否在gt里面(①)。有些prior的左上角可能在多个gt里面，这没关系，这步只关注prior的左上角是否落在gt里面。最后输出一个结果记录数组——valid_mask；

   注： prior的个数就是预测bbox的个数，所以筛选prior其实就是筛选预测bbox。因为nanodet系列是anchor-free模型，每个预测bbox坐标都是根据对应prior左上角坐标计算得来。

prior_center = priors[:, :2]  # image 坐标系下，cell 左上角坐标, shape = (num_priors, 2)
# 如果结果全为正，说明 cell 的左上角在 gt 里面
lt_ = prior_center[:, None] - gt_bboxes[:, :2]  # 所有cell左上角 与 gt左上角 的差值, shape = (num_priors, num_gt, 2)
rb_ = gt_bboxes[:, 2:] - prior_center[:, None]  # gt右下角 与 所有cell左上角 的差值, shape = (num_priors, num_gt, 2)

deltas = torch.cat([lt_, rb_], dim=-1)  # (num_priors, num_gt, 4)，坐标差值 (delta_x1, delta_y1, delta_x2. delta_y2)
# 判断 每个cell左上角 是否在 gt 里面。 先挑出4个坐标差值的最小值，再看其是否大于0，如果是置为 True.
is_in_gts = deltas.min(dim=-1).values > 0  # shape = (num_priors, num_gt)
# 如果 cell的左上角 至少在一个 gt 里面，则为 True ，否则为 False.
valid_mask = is_in_gts.sum(dim=1) > 0  # shape = (num_priors, )，得到 每个cell的左上角 与 所有gt 的关系数组

2. 用valid_mask数组初步筛选部分预测的bbox及置信度，得到初步筛选后的预测结果——valid_decoded_bbox、valid_pred_scores；

valid_decoded_bbox = decoded_bboxes[valid_mask]  # 筛选，shape = (num_valid, 4)
valid_pred_scores = pred_scores[valid_mask]  # 筛选, shape = (num_valid, num_classes)
num_valid = valid_decoded_bbox.size(0)  # 得到符合条件的 bbox 个数

3. 计算初步筛选后的预测bbox结果与gt的cost矩阵(cost_matrix)——iou_cost + cls_cost。此处的cls loss还是沿用了nanodet里面gfl的思路，即分类损失的标签是预测iou值，具体见我之前的博客，此处不做赘叙了。

# shape = (num_valid, num_gt). IOU越大，匹配效果越好，我们需要 IOU 大的 bbox结果。
pairwise_ious = bbox_overlaps(valid_decoded_bbox, gt_bboxes)  # 计算符合条件 bbox 与 gt 的 IOU值
# 转为 IOU 损失，IOU越大（靠近1），损失越小
iou_cost = -torch.log(pairwise_ious + 1e-7)

gt_onehot_label = (
    F.one_hot(gt_labels.to(torch.int64), pred_scores.shape[-1])  # shape = (num_gts, num_classes)
    .float()
    .unsqueeze(0)  # shape = (1, num_gts, num_classes)
    .repeat(num_valid, 1, 1)  # shape = (num_valid, num_gts, num_classes)
)
# shape 变为 (num_valid, num_gt, num_classes)
valid_pred_scores = valid_pred_scores.unsqueeze(1).repeat(1, num_gt, 1)
# 沿用了 gfl 的思路，用 IOU 值做分类的 label
soft_label = gt_onehot_label * pairwise_ious[..., None]
scale_factor = soft_label - valid_pred_scores
# 还是 gfl 的思路
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)
# shape = (num_valid, num_gt)。这个cost数组是分类损失与bbox损失的综合损失。 self.iou_factor = 3
cost_matrix = cls_cost + iou_cost * self.iou_factor  # IOU更重视，毕竟当前是标签分配阶段，IOU越大，标签与bbox越匹配

4. 对pairwise_ious矩阵（初步筛选后的预测bbox 与 每个gt的iou值）按列进行topk排序处理(②)，降序输出 每个gt与所有候选bbox的 前topk 个 iou值。然后按列求和，再把求和的结果直接取整，得到了每个gt可以配对的bbox个数（限制最少为1个），也就是动态k值(dynamic_k)。
   在得到每个gt的dynamic_k后，需要对 初步筛选后的预测bbox结果与gt的cost矩阵（上文的cost_matrix，下面代码里的cost）再做一次topk排序处理(③)，选出每个gt待匹配的dynamic_k个bbox。

 # select candidate topk ious for dynamic-k calculation
 candidate_topk = min(self.topk, pairwise_ious.size(0))  # 两个数之间选个最小值，免得报错
 # 降序输出 每个gt与所有候选bbox的 前topk 个 IOU值。 shape = (candidate_topk, num_gt)
 topk_ious, _ = torch.topk(pairwise_ious, candidate_topk, dim=0)
 # calculate dynamic k for each gt. 先得到每个 gt 的前topk个IOU值之和，再取整，最后做截断。得到每个gt IOU之和的整数部分
 # shape = (num_gt, ) 这个数组的每个元素是对应gt可以与几个bbox做匹配，最小值为1是因为gt肯定至少有一个bbox与之匹配
 dynamic_ks = torch.clamp(topk_ious.sum(0).int(), min=1)
 for gt_idx in range(num_gt):
     _, pos_idx = torch.topk(
         cost[:, gt_idx], k=dynamic_ks[gt_idx].item(), largest=False
     )  # 升序 动态K，选出损失最小的前 dynamic_k 个 bbox
     matching_matrix[:, gt_idx][pos_idx] = 1.0  # gt 与哪个bbox匹配，元素值置为1

下图就是对pairwise_ious按列处理的示意图。后面对cost矩阵的处理也与之类似。

5. 下面的操作就是去重，因为每个gt可以同时与多个bbox匹配，而每个bbox同时只能与一个gt匹配。

prior_match_gt_mask = matching_matrix.sum(1) > 1  # 大于 1 说明存在某些 bbox 会与多个 gt 匹配。
if prior_match_gt_mask.sum() > 0:  # 判断是否有 bbox 匹配到 多个gt 的情况
    # 下面几行的作用是 去除匹配多个 gt 的 bbox 情况，每个 bbox 只匹配一个 gt
    cost_min, cost_argmin = torch.min(cost[prior_match_gt_mask, :], dim=1)  # 选择损失最小的那个 gt 与 bbox 做匹配
    matching_matrix[prior_match_gt_mask, :] *= 0.0
    matching_matrix[prior_match_gt_mask, cost_argmin] = 1.0  # 除损失最小的gt外，其他都置为 0

具体操作为cost矩阵按行输出最小的一个gt(④)，那个gt就是当前bbox要匹配的gt，这样就避免了一个bbox匹配到多个gt的情况。

值得注意的是，虽然prior在上面步骤里出现的频率很低，更多时候都是bbox的身影。但还是要说句，正负样本定义算法的处理对象是prior而不是bbox。只是prior与bbox的数目及位置排布一致，而且bbox在上面计算里更方便，才使用bbox的。上面输出的结果最终还是会用在prior上。

最后总结一下，总共做了几次筛选（其实我上面已经有标注了），方式是什么：
① 判断prior的左上角是否在gt里面；
② 对pairwise_ious矩阵按列进行topk排序处理，得到每个gt的动态k值(dynamic_k)；
③ 对cost矩阵按列做一次topk排序处理，选出每个gt待匹配的dynamic_k个bbox；
④ cost矩阵按行输出最小的一个gt，那个gt就是每个bbox要匹配的gt。

三、疑问与个人理解

1. 动态k(dynamic_k)的计算方式；
   dynamic_k怎么来的，上面有，此处不赘叙了。首先说明，用iou来评估bbox与gt的匹配程度是正确的，但是作者把前topk个iou值加和然后取整，让我很是不解。
   试想一下，如果先把iou值与一个iou阈值做筛选，再把符合条件的bbox个数统计出来，听起来像yolo v3。但是这会引用一个先验参数——iou阈值，而先验知识是作者想避免的，所以作者没选这种方式。
   作者把iou值直接加和，获得了iou值的整体情况，整体的匹配效果越好，iou值的和就越大，最后dynamic_k也就越大，相反亦然。
2. valid_mask的赋值更新；
   这个看完代码注释会更明白点，主要是valid_mask数组在dynamic_k_matching()函数里有过赋值更新，但是函数输出时并没有输出新的valid_mask。当时猜测是内存地址没变，所以不用直接输出。查了python对象的内存机制这篇博客后，才确认是这个原因。
   总结一下就是，仅仅是对数组中某些元素重新赋值是不会新建一个对象的。但如果数组有加减乘除之类的计算，就会新建一个对象。

   valid_mask[valid_mask.clone()] = fg_mask_inboxes  #此处的赋值不会新建一块内存，所以此处的 valid_mask 与 114 行的一样
   

3. 为什么prior的左上角要在gt里面；
   之前nanodet博客里有提过，因为nanodet-plus是学习prior左上角到gt四条边界的距离，希望学习出来的四个距离都是正数，如果prior左上角在gt外部，那模型是学习不出来一个负数给它拉回去的。其次，prior左上角在gt内部，意味着有更多的特征来学习。