决定将哪个gt分配给哪个anchor的问题是目标检测中的一个重要问题,即标签分配问题。OTA提供了一个新的思考该问题的角度:把标签分配问题视为最优传输问题。
Label Assignment
以往的标签分配策略主要分为:fixed和dynamic两种。
对于fixed label assignment,通常Anchor-based的检测模型采用IoU阈值作为分配依据(如Fast-RCNN、Faster-RCNN、YoLo3、RetinaNet等),而Anchor-free检测模型有一部分将目标中心位置一定范围内的锚点指定为正样本(如FCOS、Foveabox等)。固定的分配方式通常都会采用一个固定的分配策略(例如固定的IoU阈值或固定的范围半径),对于不同尺寸、形状等的目标都是一样的,这样模型可能会得到一个次优的分配结果。
对于Dynamic label assignment,一般每个锚点的预测置信度是一个动态分配的指标,高置信度的锚点会被分配给相关的目标对象,但是动态分配的方法依然会有因为不能利用全局信息而导致锚点分配模糊的问题。
OTA利用了全局信息,将标签分配问题看为最优传输问题,主要目的是找到图像中目标全局最优的匹配。
Optimal Transport
标准的最优传输问题如下:
有\(m\)个供应者和\(n\)个需求者,第\(i\)个供应者能提供\(s_i\)件货物,而第\(j\)个需求者需要\(d_j\)件货物。每件货物从供应者\(i\)运输到需求者\(j\)所需的付出定义为\(c_{ij}\)。OT问题的目标是找到一个运输策略\(\pi^\ast=\{\pi_{i,j}\vert i=1,2,\dots,m, j=1,2,\dots,n\}\),使得运输全部货物所需的付出最小,表示如下:
\[ \begin{aligned} \min_{\pi}& \sum_{i=1}^m\sum_{j=1}^n c_{ij}\pi_{ij} \\ s.t. & \sum_{i=1}^m\pi_{ij} = d_j,\quad \sum_{j=1}^n\pi_{ij} = s_i \\ & \sum_{i=1}^ms_i = \sum_{j=1}^nd_j \\ & \pi_{ij} \ge 0,\quad i=1,2,\dots,m,\quad j=1,2,\dots,n \end{aligned} \]
该问题可以在线性时间内解决,但考虑到目标检测上涉及到特征图维度、anchor的尺度,作者使用了Sinkhorn Iteration来解决计算时间过长的问题。
Sinkhorn Iteration将上面的线性规划问题转换为如下的增加熵的正则化项的凸非线性形式:
\[ \min_\pi \sum_{i=1}^m\sum_{j=1}^n c_{ij}\pi_{ij}+\gamma E(\pi_{ij}) \]
其中\(E(\pi_{ij}) = \pi_{ij}(\log\pi_{ij}+1)\),\(\gamma\)是控制正则化强度的超参数。由拉格朗日乘数法(Lagrange Multiplier Method)有:
\[ \min_\pi \sum_{i=1}^m\sum_{j=1}^n c_{ij}\pi_{ij}+\gamma E(\pi_{ij})+\alpha_j(\sum_{i=1}^m\pi_{ij}-d_j)+\beta_i(\sum_{j=1}^n\pi_{ij}-s_i) \]
其中\(\alpha_j(j=1,2,\dots,n)\)和\(\beta_i(i=1,2,\dots,m)\)是拉格朗日乘子。对上式求导取0可得:
\[ \pi_{ij}^\ast = \exp(-\frac{\alpha_j}{\gamma})\exp(-\frac{c_{ij}}{\gamma})\exp(-\frac{\beta_i}{\gamma}) \]
令\(u_j = \exp(-\frac{\alpha_j}{\gamma})\),\(v_i = \exp(-\frac{\beta_i}{\gamma})\),\(M_{ij} = \exp(-\frac{c_{ij}}{\gamma})\)。则可将约束条件改写为:
\[ \begin{aligned} \sum_i\pi_{ij} &= u_j(\sum_iM_{ij}v_i) = d_j \\ \sum_j\pi_{ij} &= (\sum_ju_jM_{ij})v_i = s_i \end{aligned} \]
以上两个式子必须同时满足,可得以下的迭代更新式:
\[ \begin{aligned} u^{t+1}_j &= \frac{d_j}{\sum_iM_{ij}v^t_i} \\ v^{t+1}_i &= \frac{s_i}{\sum_jM_{ij}u^{t+1}_j} \end{aligned} \]
迭代完成\(T\)轮后,由下式来计算得到最佳策略\(\pi^\ast\):
\[ \pi^\ast = diag(v)Mdiag(u) \]
作者在论文中设置\(\gamma=0.1\),\(T=50\)。
Optimal Transport for Label Assignment
标签分配问题与OT问题的结合。作者将gt视为供应者,每个gt供应\(k\)个label(\(s_i=k,i=1,2,\dots,m\)),将anchor视为需求者,每个anchor需要1个label(即\(d_j=1,j=1,2,\dots,n\))。1个正标签从gt到anchor的传输代价定义为\(c^{fg}\),采用\(cls\)和\(reg\)损失的和:
\[ c^{fg}_{ij} = L_{cls}(P^{cls}_{ij}(\theta),G^{cls}_i)+\alpha L_{reg}(P^{box}_j(\theta),G^{box}_i) \]
其中\(\theta\)表示模型的参数,\(P^{cls}_j\)和\(P^{box}_j\)分别表示锚框\(a_j\)产生的预测\(cls\)分数和框的坐标,\(G^{cls}_i\)和\(G^{box}_i\)分别表示真实标注\(gt_i\)的类别和框的坐标。\(L_{cls}\)和\(L_{reg}\)为分类损失(可以采用交叉熵损失或Focal loss)和框回归损失(可以采用IoU loss、GIoU loss或SmoothL1 loss)。\(\alpha\)是平衡因子。
除了正样本的匹配,还有相当一部分anchor在训练中会被指定为负样本,所以作者还增加了负样本的分配,引入了背景供应者(只提供负标签)。在标准OT问题中,供应者的总供应量应该与需求者的总需求量相同,所以背景供应量计算如下:
\[ c^{bg}_j =L_{cls}(P^{cls}_j(\theta), \varnothing) \]
其中\(\varnothing\)表示背景类。将\(c^{bg}\in R^{1\times n}\)添加到\(c^{fg}\in R^{m\times n}\)末尾获得完整的代价矩阵\(c\in R^{(m+1)\times n}\),供应向量\(s\)相应地更新为下式:
\[ s_i = \begin{cases} k, & if\quad i\le m \\ n-m\times k, & if\quad i=m+1 \end{cases} \]
已知代价矩阵\(c\)、供应向量\(s\in R^{m+1}\)和需求向量\(d\in R^n\),通过求解该OT问题即可得到最优的分配策略\(\pi^\ast\in R^{(m+1)\times n}\)。
一些优化设计
- Center prior。中心先验,即在目标中心指定区域内选取anchor进行分配。理论上OTA可以将全局内的任意anchor指定为正样本,而通过center prior可以让OTA更关注潜在的正样本区域,可以稳定模型前期的训练。作者采用的具体操作是对于每个gt,依据anchor与gt的中心距离在每一层FPN中选取前\(r^2\)个anchor作为正样本候选,对于其余没有入选的anchor,会在代价矩阵相应位置加上一个常数来降低anchor被选为正样本的概率。
- 动态的\(k\)估计。一般对于每个gt应该具有不同的合适的正锚框/点个数,因为每个gt的尺寸、尺度以及遮挡程度都不同。作者依据预测框与gt的IoU数值对每个gt选取前q个IoU数值,并将其求和得到的值作为该gt合适的正锚框/点个数(因为某个gt合适的正锚框/点数量应该与能够很好回归该gt的锚框/点数量正相关)。
综合以上内容,可以得到OTA的伪代码如下: