OTA: Optimal Transport Assignment for Object Detection原理与代码解读
paper：OTA: Optimal Transport Assignment for Object Detection
code：https://github.com/Megvii-BaseDetection/OTA
背景
标签分配（Label Assignment）是目标检测中重要的一环，经典的标签分配策略采用预定义的规则为每个anchor匹配对应的gt或背景类。比如RetinaNet采用IoU作为划分正负样本的阈值标准，anchor-free检测器比如FCOS将ground truth物体的bbox内或bbox中心区域内的anchor point作为正样本。这种静态分配策略忽略了这样一个事实，即对于不同大小、形状、遮挡状态的对象，最适合的正负样本划分的边界可能是不同的。
基于此很多动态分配方法被提出，比如ATSS基于统计特征为每个gt设置划分边界，Freeanchor、Autoassign、PAA等方法提出anchor的预测分数可以作为一个合适的指标用来设计动态分配策略。
但是，不考虑上下文单独的为每个gt分配正负样本的方法可能不是最优的。对于模糊的anchor，即可能作为正样本分配给多个gt的anchor，现有的策略都是基于人工定义的准则，比如Min Area或Max IoU。作者指出把ambiguous anchor分配给任一个gt，对其他gt的学习都是不利的（introduce harmful gradients w.r.t. other gts），因此分配还需要更多的信息。一个更好的分配策略应该摆脱对每个gt单独追求最优分配的思想，转而全局最优的思想，找到一张图像中所有gt的综合最优分配策略。
本文的创新点
本文提出把标签分配当做最优传输问题，具体是把每个gt定义成一个supplier，它可以提供一定数量的label。把每个anchor定义成demander，它需要一个label。如果一个anchor从某个gt那得到了足够数量的positive label，这个anchor就被当做这个gt的一个正样本。每个gt可以提供的positive label的数量可以理解为这个gt在训练过程中需要多少个正样本来更好的收敛。每对anchor-gt的传输cost定义为它们之间的分类和回归loss的加权和。此外，背景类也被定义为supplier，它提供negative label，anchor-background之间的传输cost定义为它们之间的分类loss。这样标签分配问题就被转化为了最优传输问题，最终是为了找到全局最优的分配方法而不再是为每个gt单独寻找最优anchor。
具体方法
Optimal Transport
最优传输问题可以表述为：假设有 (m) 个supplier和 (n) 个demander，第 (i) 个supplier有 (s_{i}) 个物品，第 (j) 个demander需要 (d_{j}) 个物品，每个物品从第 (i) 个supplier运到第 (j) 个demander的运输运输成本为 (c_{ij})，最优传输的目标是找到一个最优传输方案 (pi^{*}=left { pi_{i,j}|i=1,2,...m,j=1,2,...n right } ) 能以最小的运输成本把所有的物品从supplier运输到demander。
OT for Label Assignment
对于目标检测问题，假设一张图片有 (m) 个gt和 (n) 个anchor（所有FPN level加起来），每个gt当做一个supplier，持有 (k) 个正标签 ((i.e.,s_{i}=k,i=1,2,...,m))，每个anchor当做一个demander，需要一个标签 ((i.e.,d_{j}=1,j=1,2,...,n))。从 (gt_{i}) 传输一个正标签到anchor (a_{j}) 的运输成本 (f^{fg}) 定义为它们之间的分类损失和回归损失的加权和
其中 (theta) 是模型参数，(P_{j}^{cls}) 和 (P_{j}^{reg}) 分别表示anchor (a_{j}) 的预测的分类得分和bounding box。(G_{i}^{cls}) 和 (G_{i}^{box}) 分别表示 (gt_{i}) 的ground truth类别和bounding box。(L_{cls}) 和 (L_{reg}) 分别表示交叉熵loss和IoU loss，也可以分别替换成Focal loss和GIoU/Smooth L1 loss，(alpha) 是权重系数。
此外，还有另一种提供负标签的supplier，背景类。在标准的最优传输问题中，supply的数量和demand的数量是相等的。因此背景类一共可以提供 (n-mtimes k) 个负标签，从背景类传输一个负标签到 (a_{j}) 的成本为
其中 (oslash) 表示背景类，把 (c^{bg}in mathbb{R}^{1times n}) 拼接到 (c^{fg}in mathbb{R}^{mtimes n}) 的最后一行即得到了完整的cost matrix (cin mathbb{R}^{(m+1)times n})。supply vector (s) 需要按下式更新
现在有了cost matrix (c)，supply vector (sin mathbb{R}^{m+1})，demand vector (din mathbb{R}^{n})，则最优传输路径 (pi^{*}in mathbb{R}^{(m+1)times n}) 可通过现有的Sinkhorn-Knopp Iteration算法求得。得到 (pi^{*}) 后，对应的标签分配就是将每个anchor分配给传输给这个anchor最多标签的gt。
Advanced Designs
Center Prior
center prior即只从gt的中心有限区域挑选正样本，而不是整个bounding box范围内选择。强迫模型关注潜在positive areas即中心区域有助于稳定训练，特别是在训练的早期阶段，模型的最终性能也会更好。作者发现center prior对OTA的训练也有帮助，因此引入了center prior策略。
具体做法是，对于每个gt，只挑选每个FPN层中距离bounding box中心最近的 (r^{2}) 个anchor，对于bounding box内 (r^{2}) 之外的anchor，cost matrix中对应的cost会加上一个额外的常数项cost，这样就减少了训练阶段它们被分配为正样本的概率。
Dynamic (k) Estimation
每个gt需要的正样本数量应该是不同的并且基于很多因素，比如物体大小、尺度、遮挡情况等。由于很难将这些因素和所需anchor数量直接映射起来，本文提出了一种简单有效的方法，根据预测框和对应gt的IoU值来粗略估计每个gt合适的正样本数量。具体来说，对于每个gt，选择IoU最大的 (q) 个个预测，将这 (q) 个IoU值的和作为这个gt正样本数量的粗略估计值。这样做是基于直觉：某个gt的所需合适的postive anchor数量与和这个gt拟合的很好的anchor的数量正相关。
OTA的完整流程如下图所示
包含center prior和dynamic k estimation的完整流程伪代码如下所示
代码解读
这里batch_size=2，输入shape=(2, 3, 1085, 800)，前景loss权重系数 (alpha=1.5)，center prior超参 (r=2.5)，dynamic (k) estmation中 (q=20)。
其中line96计算前景loss和中的 1e6*(1-is_in_boxes.float()) 就是中心区域外的anchor额外加的常数项cost，line105将背景的cost拼接到前景cost矩阵最后就得到了最终的cost matrix，这里的loss就是cost matrix。mu和nu分别是上面的supply vector (s) 和 demand vector (d)。
核心代码如下，加了一些注释，其中sinkhorn算法没有专门了解原理，这里就直接用吧。
def get_ground_truth(self, shifts, targets, box_cls, box_delta, box_iou):
# shifts
# [[(13600,2),(3400,2),(850,2),(221,2),(63,2)],
# [(13600,2),(3400,2),(850,2),(221,2),(63,2)]]
# targets
# [Instances(num_instances=2, image_height=1085, image_width=800,
# fields=[gt_boxes = Boxes(tensor([[216.9492, 217.0000, 605.6497, 965.1979], [246.3277, 160.4896, 501.6949, 641.9583]], device='cuda:0')),
# gt_classes = tensor([12, 14], device='cuda:0'), ]),
# Instances(num_instances=2, image_height=1085, image_width=800,
# fields=[gt_boxes = Boxes(tensor([[216.9492, 217.0000, 605.6497, 965.1979], [246.3277, 160.4896, 501.6949, 641.9583]], device='cuda:0')),
# gt_classes = tensor([12, 14], device='cuda:0'), ])]
gt_classes = []
gt_shifts_deltas = []
gt_ious = []
assigned_units = []
box_cls = [permute_to_N_HWA_K(x, self.num_classes) for x in box_cls]
# [(2,13600,20),(2,3400,20),(2,850,20),(2,221,20),(2,63,20)]
box_delta = [permute_to_N_HWA_K(x, 4) for x in box_delta]
# [(2,13600,4),(2,3400,4),(2,850,4),(2,221,4),(2,63,4)]
box_iou = [permute_to_N_HWA_K(x, 1) for x in box_iou]
# [(2,13600,1),(2,3400,1),(2,850,1),(2,221,1),(2,63,1)]
box_cls = torch.cat(box_cls, dim=1) # (2,18134,20)
box_delta = torch.cat(box_delta, dim=1) # (2,18134,4)
box_iou = torch.cat(box_iou, dim=1) # (2,18134,1)
for shifts_per_image, targets_per_image, box_cls_per_image,
box_delta_per_image, box_iou_per_image in zip(
shifts, targets, box_cls, box_delta, box_iou):
shifts_over_all = torch.cat(shifts_per_image, dim=0) # (18134,2)
gt_boxes = targets_per_image.gt_boxes # (2,4)
# In gt box and center.
deltas = self.shift2box_transform.get_deltas(
shifts_over_all, gt_boxes.tensor.unsqueeze(1)) # (18134,2),(2,1,4) -> (2,18134,4)
is_in_boxes = deltas.min(dim=-1).values > 0.01 # (2,18134)
center_sampling_radius = 2.5
centers = gt_boxes.get_centers() # (2,2),
# tensor([[388.7006, 591.0990],
# [425.9887, 401.2239]], device='cuda:0')
# 因为数据增强的, gt_bboxes和centers每次运行结果都会变化
is_in_centers = []
for stride, shifts_i in zip(self.fpn_strides, shifts_per_image): # [8, 16, 32, 64, 128], _
radius = stride * center_sampling_radius
center_boxes = torch.cat((
torch.max(centers - radius, gt_boxes.tensor[:, :2]),
torch.min(centers + radius, gt_boxes.tensor[:, 2:]),
), dim=-1) # (2,4)
center_deltas = self.shift2box_transform.get_deltas(
shifts_i, center_boxes.unsqueeze(1)) # (13600,2),(2,1,4) -> (2,13600,4)
is_in_centers.append(center_deltas.min(dim=-1).values > 0)
is_in_centers = torch.cat(is_in_centers, dim=1) # (2,18134)
del centers, center_boxes, deltas, center_deltas
is_in_boxes = (is_in_boxes & is_in_centers)
num_gt = len(targets_per_image)
num_anchor = len(shifts_over_all)
shape = (num_gt, num_anchor, -1) # (2,18134,-1)
gt_cls_per_image = F.one_hot(
targets_per_image.gt_classes, self.num_classes
).float() # (2,20)
with torch.no_grad():
loss_cls = sigmoid_focal_loss_jit(
box_cls_per_image.unsqueeze(0).expand(shape), # (18134,20)->(1,18134,20)->(2,18134,20)
gt_cls_per_image.unsqueeze(1).expand(shape), # (2,20)->(2,1,20)->(2,18134,20)
alpha=self.focal_loss_alpha, # 0.25
gamma=self.focal_loss_gamma, # 2
).sum(dim=-1) # (2,18134,20)->(2,18134)
loss_cls_bg = sigmoid_focal_loss_jit(
box_cls_per_image, # (18134,20)
torch.zeros_like(box_cls_per_image),
alpha=self.focal_loss_alpha,
gamma=self.focal_loss_gamma,
).sum(dim=-1) # (18134,20)->(18134)
gt_delta_per_image = self.shift2box_transform.get_deltas(
shifts_over_all, gt_boxes.tensor.unsqueeze(1) # (18134,2), (2,4)->(2,1,4)
) # (2,18134,4)
ious, loss_delta = get_ious_and_iou_loss(
box_delta_per_image.unsqueeze(0).expand(shape), # (18134,4)->(1,18134,4)->(2,18134,4)
gt_delta_per_image,
box_mode="ltrb",
loss_type='iou'
) # (2,18134),(2,18134)
loss = loss_cls + self.reg_weight * loss_delta + 1e6 * (1 - is_in_boxes.float()) # 1.5
# (2,18134)
# Performing Dynamic k Estimation
topk_ious, _ = torch.topk(ious * is_in_boxes.float(), self.top_candidates, dim=1) # (2,18134),20 -> (2,20)
mu = ious.new_ones(num_gt + 1) # torch.Size([3]), tensor([1., 1., 1.], device='cuda:0')
mu[:-1] = torch.clamp(topk_ious.sum(1).int(), min=1).float() # s_{i}(i=1,...,m)
mu[-1] = num_anchor - mu[:-1].sum() # s_{m+1}
nu = ious.new_ones(num_anchor) # (18134), d_{j}(j=1,..,n)
loss = torch.cat([loss, loss_cls_bg.unsqueeze(0)], dim=0) # (2,18134),(18134)->(1,18134), -> (3,18134)
# Solving Optimal-Transportation-Plan pi via Sinkhorn-Iteration.
_, pi = self.sinkhorn(mu, nu, loss) # (3,),(18134,),(3,18134) -> (3,18134)
# Rescale pi so that the max pi for each gt equals to 1.
rescale_factor, _ = pi.max(dim=1) # (3,)
pi = pi / rescale_factor.unsqueeze(1) # (3,18134)
max_assigned_units, matched_gt_inds = torch.max(pi, dim=0)
gt_classes_i = targets_per_image.gt_classes.new_ones(num_anchor) * self.num_classes
fg_mask = matched_gt_inds != num_gt
gt_classes_i[fg_mask] = targets_per_image.gt_classes[matched_gt_inds[fg_mask]]
gt_classes.append(gt_classes_i)
assigned_units.append(max_assigned_units)
box_target_per_image = gt_delta_per_image.new_zeros((num_anchor, 4))
box_target_per_image[fg_mask] =
gt_delta_per_image[matched_gt_inds[fg_mask], torch.arange(num_anchor)[fg_mask]]
gt_shifts_deltas.append(box_target_per_image)
gt_ious_per_image = ious.new_zeros((num_anchor, 1))
gt_ious_per_image[fg_mask] = ious[matched_gt_inds[fg_mask],
torch.arange(num_anchor)[fg_mask]].unsqueeze(1)
gt_ious.append(gt_ious_per_image)
return torch.cat(gt_classes), torch.cat(gt_shifts_deltas), torch.cat(gt_ious)Experiments
Alation Studies and Analysis
Effects of Individual Components
OTA可以既可以用于anchor-based detector也可以用于anchor-free detector，本文采用FCOS，同时额外加入了IoU分支，从下图可以看出随着添加IoU branch、center prior、dynamic k estimation，性能持续提升，并且比对应的原始FCOS的精度要高。
Effects of (r)
center prior的半径 (r) 控制每个gt的正样本数量，(r) 值小，只有最靠近gt中心的高质量anchor才被当做正样本，有助于模型的学习。(r) 越大，引入的低质量的正样本anchor越多，导致了优化过程中潜在的不稳定。从下表可以看出，随着 (r) 的增大，三种模型的精度都出现了不同程度的下降，但OTA下降的最少，表明OTA对 (r) 值的变化不那么敏感，同时不同的 (r) 值下，OTA的精度也是最高的。
Ambiguous Anchors Handling
当发生遮挡或者多个对象靠的非常近时，一个anchor可能是多个ground truth的合格候选对象（比如Faster RCNN中一个anchor与多个gt的IoU都大于0.5），这种anchor定义为ambiguous anchor。之前的方法主要通过人工设定的规则来处理这种情况，比如Min Area、Max IoU、Min Loss等。本文将 (max pi^{*}_{j}<0.9) 的anchor (a_{j}) 定义为ambiguous anchor，然后统计在不同的 (r) 值下ATSS、PAA、OTA的ambiguous anchor的数量以及对应的精度。从上表(2)中可以看出，随着 (r) 的增大，ATSS中ambiguous anchor的数量显著增加，AP也降了1.8个点。PAA中ambiguous anchor的数量对 (r) 的变化不那么敏感，但AP也降了0.8个点。而OTA中ambiguous anchor的数量既对 (r) 的变化不敏感，和ATSS、PAA相比数量也是最少的，同时AP也只下降了0.3个点。这是因为当多个gt试图将positive label传输到同一个anchor时，OT算法会基于全局最小传输成本的准则自动解决它们之间的冲突。
Effects of (k)
如下表所示，作者对比了 (k) 设置为不同的常数值以及采用dynamic (k) 时模型的精度，可以看出随着 (k) 的增大，模型精度越来越高，当 (k) 取10或12时，模型达到最高的精度，随后开始下降。但最高的精度也比采用dynamic (k) 的精度低。从直觉上讲，每个gt的大小、尺度、遮挡情况都不同，因此每个gt所需的postive anchor的数量应该也是不同的。
Comparison with State-of-the-art Methods
从下表可以看出，采用ResNet-101-FPN结构，OTA的AP达到了45.3%，超过了其它所有相同backbone的方法，如ATSS（43.6% AP）、AutoAssign（44.5% AP）、PAA（44.6% AP）。