【NIPS2023】Rank-DETR for High Quality Object Detection

【NIPS2023】Rank-DETR for High Quality Object Detection

机构：清华大学、北京大学、剑桥大学、微软亚洲研究院

论文地址：https://arxiv.org/abs/2310.08854

代码地址：https://github.com/LeapLabTHU/Rank-DETR

作者简介：黄高，清华大学博士学位，康奈尔大学计算机系博士后，清华大学自动化系助理教授、博士生导师，获阿里巴巴“达摩院青橙奖”、2019年吴文俊人工智能优秀青年奖等。代表作DenseNet获得CVPR2017年最佳论文、Stochastic Depth。研究方向包括动态神经网络、高效深度学习。

本文考虑到DETR模型中query的重要性存在差异，致力于改进高IoU情况下（例如AP@75）的检测性能，首次提出基于排序思想的Rank-DETR，在Transformer中引入排序相关的网络层、排序导向的损失函数和匈牙利匹配损失。在COCO数据集上的性能高于DINO、Align-DETR、GroupDETR等baseline，与Stable-DINO、MS-DETR、Salience-DETR相当，弱于DDQ-DETR、Co-DETR、Relation-DETR等SOTA方法。

文章贡献/创新点

• 在Transformer Decoder中提出了基于rank机制改进的分类头和query排序层。
• 在损失函数（网络损失和匈牙利匹配损失）中对分类和回归分支进行对齐，使得高置信度的query也具有高IoU。
• 实验验证了所提方法的有效性，并将rank机制引入到已有DETR方法中验证了有效性。

排序相关的结构设计

排序自适应的分类头

常规的DETR方法中，backbone提取多尺度特征，transformer将其映射为6层Decoder输出（两阶段方法还会多1层Encoder输出），每层的输出都是$n$个query（原始DETR中$n=100$、DeformableDETR中$n=300$、DINO中$n=900$），针对第$l$层的每个query，表示为$q_i^l$，head将其映射为分类结果${\mathbf{p}}_i^l$+回归结果，其中分类头是单层全连接：

$${\mathbf{p}}_i^l=\mathrm{Sigmoid}({\mathbf{r}}_i^l),{\mathbf{t}}_i^l={\text{MLP}}_{\text{cls}}({\mathbf{q}}_i^l)$$

本文提出的排序自适应分类头其实就是为每个query增加了对应的embedding，两者加起来再进行分类：

$${\mathbf{p}}_i^l=\mathrm{Sigmoid}({\mathbf{t}}_i^l+{\mathbf{s}}_i^l),{\mathbf{t}}_i^l={\text{MLP}}_{\text{cls}}({\mathbf{q}}_i^l)$$

这里的${\mathbf{s}}_i^l$表示第$l$层第$i$个query的embedding。所有的$\mathbf{s}$都作为网络参数自适应学习。

Query排序层

Transformer Decoder输入的query包含两部分，分别是content query和position query。其中content query作为网络可以学习的参数进行初始化，position query来自Transformer Encoder输出的候选框。RankDETR在每个Transformer解码层后增加了一个query rank layer来对这两种query进行排序。

对内容query进行排序

排序依据是每层输出的分类结果${\hat{\mathcal{P}}}^{l-1}={\mathrm{MLP}}_{\text{cls}}({\mathcal{Q}}_c^{l-1})$，排序层会按照该置信度对内容query${\mathcal{Q}}_c$进行降序排序，排序后的${\hat{\mathcal{Q}}}_c^l$与${\mathcal{C}}^l$进行拼接，这里的${\mathcal{C}}^l$同样是作为网络参数自适应学习。拼接后的结果在channel维度就变成了2倍，因此再经过全连接进行降维，得到下一层的query。整个流程：

$${\overline{\mathcal{Q}}}_c^l-{\mathrm{MLP}}_{\text{fuse}}({\hat{\mathcal{Q}}}_c^{l-1}||{\mathcal{C}}^l),{\hat{\mathcal{Q}}}_c^{l-1}=\mathrm{Sort}({\mathcal{Q}}_c^l-1;{\hat{\mathcal{P}}}^{l-1})$$

对位置query进行排序

位置query的排序方式在各个DETR方法中有所不同，针对H-DETR和Deformable DETR，作者也是按照${\hat{\mathcal{P}}}^{l-1}$对每一层位置query进行降序排序：

$${\overline{\mathcal{Q}}}_p^l=\mathrm{Sort}({\overline{\mathcal{Q}}}_p^{l-1};{\hat{\mathcal{P}}}^{l-1})$$

由于DINO-DETR的位置query是上一层的位置query经过bounding box微调的结果，作者没有直接对${\hat{\mathcal{Q}}}_p^l$进行排序，而是先对检测框进行排序，再由排序后的检测框生成下一层的query：

$${\overline{\mathcal{Q}}}_p^l=\mathrm{PE}({\overline{\mathcal{B}}}^{l-1}),{\overline{B}}^{l-1}=\mathrm{Sort}({\mathcal{B}}^{l-1};{\hat{\mathcal{P}}}^{l-1})$$

其中$\mathrm{PE}$表示正余弦编码和多层全连接。

代码实现

作者主要是基于H-DETR实现的Rank-DETR，因此代码没给出DINO-DETR的排序方式。

# 对内容进行排序
output = torch.gather(
    output, 1, rank_indices.unsqueeze(-1).repeat(1, 1, output.shape[-1])
)
# 排序后与C进行拼接，然后经过MLP
concat_term = self.pre_racq_trans[layer_idx - 1](
    self.rank_aware_content_query[layer_idx - 1].weight[:output.shape[1]].unsqueeze(0).expand(output.shape[0], -1, -1)
)
output = torch.cat((output, concat_term), dim=2)
output = self.post_racq_trans[layer_idx - 1](output)

# 对未知进行排序
query_pos = torch.gather(
    query_pos, 1, rank_indices.unsqueeze(-1).repeat(1, 1, query_pos.shape[-1])
)


# 省略中间代码......

# 获得排序依据：训练时有one2one query和one2many query，要分别对排序，推理时只有one2one query
if self.training:
    rank_indices_one2one  = torch.argsort(rank_basis[:, : self.num_queries_one2one], dim=1, descending=True) # tensor shape: [bs, num_queries_one2one]
    rank_indices_one2many = torch.argsort(rank_basis[:, self.num_queries_one2one :], dim=1, descending=True) # tensor shape: [bs, num_queries_one2many]
    rank_indices = torch.cat(
        (
            rank_indices_one2one,
            rank_indices_one2many + torch.ones_like(rank_indices_one2many) * self.num_queries_one2one
        ),
        dim=1,
    ) # tensor shape: [bs, num_queries_one2one+num_queries_one2many]
else:
    rank_indices = torch.argsort(rank_basis[:, : self.num_queries_one2one], dim=1, descending=True)

排序相关的损失设计

损失函数设计

一般DETR的损失包括三部分，分类损失、定位损失和GIoU损失：

$$-{\lambda }_1\mathrm{GIoU}(\hat{\mathbf{b}},\mathbf{b})+{\lambda }_2{\ell }_1(\hat{\mathbf{b}},\mathbf{b})+{\lambda }_3\mathrm{FL}(\hat{\mathbf{p}}[c])$$

作者提出的改进就在于其中的$\mathrm{FL}$分类损失上，作者将分类目标从原始的二分类0-1目标替换为了基于$\mathrm{IoU}$的分类目标：

$${\mathrm{FL}}^{\text{GIoU}}(\hat{\mathbf{p}}[c])=-|t-\hat{\mathbf{p}}[c]{|}^{\gamma }\cdot [t\cdot \log (\hat{\mathbf{p}})]+(1-t)\cdot \log (1-\hat{\mathbf{p}}[c])$$

其中$t=(\mathrm{GIoU}(\hat{\mathbf{b}},\mathbf{b})+1)/2$。文章还对比了VFL的损失函数：

$$\mathrm{VFL}(\hat{\mathbf{p}}[c])=-t\cdot [t\cdot \log (\hat{\mathbf{p}}[c])+(1-t)\cdot \log (1-\hat{\mathbf{p}}[c])]$$

可以看到，两者区别就在于对正样本监督的权值从$t$变成了$|t-\hat{\mathbf{p}}[c]{|}^{\gamma }$。

思考：一般来说DETR正样本监督式很稀缺的，因此在设计损失函数的时候（例如VFL和TOOD损失），正样本通常不进行难例挖掘。这里作者对正样本同样进行了类似GFL的难例挖掘设计，可能是考虑到H-DETR本身通过one2many的匹配设计弥补了正样本监督稀缺的问题，在此基础上进行难例挖掘则可能提升性能。

匹配损失函数设计

常规的匹配损失是对分类、回归和IoU损失进行加权，加权比例通常是2、5、2，文章提出了高阶的匹配损失：

$${\mathcal{L}}_{\text{Hungarian}}^{\text{high-order}}=\hat{\mathbf{p}}[c]\cdot {\text{IoU}}^{\alpha }$$

其中$\hat{\mathbf{p}}[c]$是分类头的输出。在代码实现中，前期仍然是使用常规的匹配损失，后期则是换成作者提出的高阶损失。

实验结果

从消融实验中来看，改进损失函数和匹配损失对性能提升的效果最大：

文章另外对比了自己提出的损失和VFL，发现使用自己提出的排序损失能够达到49.8%的AP，而VFL则只有49.5%AP，说明作者提出的排序损失能够更好地建模boxes之间地距离。

Rank-DETR虽然引入了2个结构上的改进，但都比较轻量化；而损失函数的改进不会对推理性能有影响，因此总体FPS基本没有太大下降：

和主流方法的性能对比实验可以看出，Rank-DETR在1$\times$下的性能达到了50.2，高于DINO，但相比其他方法（DDQ-DETR、Relation-DETR、Co-DETR），性能还是弱一些，但胜在对推理速度没有太大影响。