【论文阅读】Transformer模型在CV目标检测任务中的应用

论文：End-to-End Object Detection with Transformers

一、 动机

作者认为，现有的目标检测技术都是基于region proposal、anchor、window center、NMS这些“间接”的技术，而不是“直接”对bounding box和类别进行预测，导致了大量的调参工作以及中间步骤（如NMS）的损失。为了解决这个问题，作者提出了一种更直接的方案，直接从CNN特征预测框的位置和类别。

二、 方法

该方法将Transformer模型引入目标检测任务中，因此方法名称为DETR（Detection with Transformers）。该方法的架构图如下图所示：

从图中可以看出，与CV领域现有的目标检测算法相比，DETR方法的backbone部分与现有方法并没有区别，只是在模型的head部分加入了Transformer元素。具体来说，模型的训练过程包含以下步骤：

• 将原始图像输入到模型的backbone（如ResNet）中，得到一个H*W*C的feature map，其中H和W分别是feature map的高度和宽度，C是feature map的通道数；
• 通过一个可学习的卷积层，使用1*1的卷积核，将feature map的通道数降低到d（d < C），得到一个H*W*d的feature map；
• 为了使feature map的维度和Transformer模型匹配，将feature map进行重新排列，得到一个d*HW维的矩阵；
• 对上述d*HW的矩阵进行位置编码；
• DETR模型中的Transformer部分以编码后的d*HW维矩阵作为encoder的输入，以N个学习到的位置编码作为decoder的输入，经过一系列运算后，decoder输出一个N*m维的矩阵，其中m表示decoder输出特征的维度；
• 利用一个小的全连接网络，基于上述N*m维矩阵进行预测，得到N个结果中每个结果的类别（前景类别或背景）、bounding box位置；
• 使用匈牙利算法将N个模型预测结果与标注结果进行匹配（标注结果补齐到N个，包含背景）；
• 计算loss的值，并进行反向传播，对模型参数进行更新。

推理阶段，流程与训练阶段基本一致，但是只执行到第（6）步，不进行后续的匹配、loss计算步骤。

作者指出，该算法在训练阶段需要进行更多次数的迭代才能收敛，但是在推理阶段与现有目标检测算法没有明显区别。

三、 实验结果

论文作者在MS COCO数据上进行了实验，实验结果如下表所示：

从表中的结果可以看出，在参数量与现有方法相近的情况下，DETR方法可以达到更高的检测精度（AP），并且精度的提升主要来自于中等尺寸物体、大尺寸物体检测效果的提升（事实上小物体的检测效果略有下降）。如果将本方法与FPN等技术结合在一起，在小目标检测上的效果也值得期待。

四、 方法扩展

论文中，作者还展示了本方法不仅仅局限于目标检测任务，在语义分割任务上也可以取得不错的效果。具体来说，通过对Transformer的输入、输出进行调整，使用相同的backbone可以进行语义分割任务，如下图所示：

从实验结果来看，虽然DETR方法在语义分割任务上的精度没有那么惊艳，但是也达到了SOTA的级别：

五、 总结

近年来，随着Transformer模型在自然语言领域的兴起，越来越多的研究工作也在尝试将其应用到其他AI相关领域。本文中的DETR算法就是Transformer在CV领域中的目标检测问题上的一个初步尝试。得益于Transformer方法无限大的感受野，引入Transformer后可以省去目标检测中的大量手工设计的操作（如NMS、region proposal等），并且从实验结果来看，DETR方法的效果也相当不错。