DETR: 使用Transformer实现端到端目标检测.

本文提出了Detection Transformer (DETR),一种将目标检测建模为集合预测任务的方法。该方法不依赖于目标检测中需要先验知识的手工设计组件,比如anchor设置和非极大值抑制;而是通过二分匹配执行唯一的预测,并构建基于集合的全局损失。给定一个可学习的目标查询集,DETR通过Transformer的编码器-解码器结构提取目标之间的关系和图像的全局上下文信息,并行地输出最终的预测集合。

1. DETR的模型框架

实现直接的集合预测需要两个要素:对预测结果和真实标签进行唯一匹配的集合预测损失,以及一次性预测一组对象并对其关系进行建模的网络结构。基于此作者设计了DETR结构:

(1) 损失函数

DETR在单次推断过程中一次性的给出$N$个目标预测(包括目标类别和边界框)。记$y$为目标的真实标签targets,\(\hat{y}=\{\hat{y}_i\}_{i=1}^{N}\)是$N$个预测结果outputs。通常预测结果数量$N$远大于图像中实际存在目标的典型数量,因此对图像标签增加空集(no object)直至达到$N$个。

""" 
Params:
    outputs: This is a dict that contains at least these entries:
         "pred_logits": Tensor of dim [batch_size, num_queries, num_classes] with the classification logits
         "pred_boxes": Tensor of dim [batch_size, num_queries, 4] with the predicted box coordinates
    targets: This is a list of targets (len(targets) = batch_size), where each target is a dict containing:
         "labels": Tensor of dim [num_target_boxes] (where num_target_boxes is the number of ground-truth
                   objects in the target) containing the class labels
         "boxes": Tensor of dim [num_target_boxes, 4] containing the target box coordinates
"""

为了构造损失函数,需要建立这两个集合之间的最优二分匹配(bipartite matching),这是通过寻找使得代价最小的一个排列$\sigma$实现的:

\[\hat{\sigma} = \mathop{\arg \min}_{\sigma \in \Sigma_N} \sum_{i}^{N} \mathcal{L}_{\text{match}}(y_i,\hat{y}_{\sigma(i)})\]

其中\(\mathcal{L}_{\text{match}}(y_i,\hat{y}_{\sigma(i)})\)是真值$y_i$和对应预测\(\hat{y}_{\sigma(i)}\)之间的匹配代价,匹配过程是通过由匈牙利算法构造的最优指派,这种匹配是一对一的,不会产生重复的匹配结果。

from scipy.optimize import linear_sum_assignment
# assert C.shape == (batch_size, num_queries, batch_size * num_target_boxes)
sizes = [len(v["boxes"]) for v in targets] # batch_size x [num_target_boxes]
indices = [linear_sum_assignment(c[i]) for i, c in enumerate(C.split(sizes, -1))]
return [(torch.as_tensor(i, dtype=torch.int64), torch.as_tensor(j, dtype=torch.int64)) for i, j in indices]

# 注:源代码中对代价矩阵的计算没有避开不同样本之间的预测和真值交互
#     而是通过split操作对第i个样本只取对应的矩阵块(存在许多不必要的计算)

匹配代价同时考虑类别预测结果和预测边界框的准确程度,真值$y_i$是由预测类别$c_i$和边界框(相对图像的)中心坐标与尺寸$b_i \in [0,1]^4$组成的。若记预测\(\hat{y}_{\sigma(i)}\)的类别概率为\(\hat{p}_{\sigma(i)}(c_i)\),则匹配代价计算为:

\[\mathcal{L}_{\text{match}}(y_i,\hat{y}_{\sigma(i)}) = -\Bbb{I}(c_i \ne \Phi) \hat{p}_{\sigma(i)}(c_i) + \Bbb{I}(c_i \ne \Phi) \mathcal{L}_{\text{box}}(b_i,\hat{b}_{\sigma(i)})\]

在匹配代价中,直接使用类别概率\(\hat{p}_{\sigma(i)}(c_i)\)代替交叉熵中的对数概率,这使得其与边界框损失具有同单位度量,并且具有更好的实际表现。边界框损失若直接使用L1损失,会使得小边界框和大边界框具有不同的尺度,因此作者采用L1损失和GIoU损失的线性组合,并根据一个批量样本内目标的数量进行规范化:

\[\mathcal{L}_{\text{box}}(b_i, \hat{b}_{\hat{\sigma}}(i)) = \lambda_{\text{iou}}\mathcal{L}_{\text{iou}}(b_i, \hat{b}_{\hat{\sigma}}(i))+\lambda_{L1} ||b_i- \hat{b}_{\hat{\sigma}}(i)||_1\]
bs, num_queries = outputs["pred_logits"].shape[:2]

# We flatten to compute the cost matrices in a batch
out_prob = outputs["pred_logits"].flatten(0, 1).softmax(-1)  # [batch_size * num_queries, num_classes]
out_bbox = outputs["pred_boxes"].flatten(0, 1)  # [batch_size * num_queries, 4]

# Also concat the target labels and boxes
tgt_ids = torch.cat([v["labels"] for v in targets])  # [batch_size * num_target_boxes]
tgt_bbox = torch.cat([v["boxes"] for v in targets])  # [batch_size * num_target_boxes, 4]

# Compute the classification cost. Contrary to the loss, we don't use the NLL,
# but approximate it in 1 - proba[target class].
# The 1 is a constant that doesn't change the matching, it can be ommitted.
cost_class = -out_prob[:, tgt_ids]  # [batch_size * num_queries]

# Compute the L1 cost between boxes
cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1) # [batch_size * num_queries, batch_size * num_target_boxes]

# Compute the giou cost betwen boxes
cost_giou = -generalized_box_iou(out_bbox, tgt_bbox) # [batch_size * num_queries, batch_size * num_target_boxes]

# Final cost matrix
C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou
C = C.view(bs, num_queries, -1).cpu() # [batch_size, num_queries, batch_size * num_target_boxes]

预测集合和真值集合匹配完成后可以计算损失函数,采用类别预测的负对数和边界框损失的线性组合:

\[\mathcal{L}_{\text{Hungarian}}(y,\hat{y}) = \sum_{i=1}^{N} [-\log \hat{p}_{\hat{\sigma}(i)}(c_i)+\Bbb{I}(c_i \ne \Phi) \mathcal{L}_{\text{box}}(b_i, \hat{b}_{\hat{\sigma}}(i)) ]\]

其中$\hat{\sigma}$表示构造的最优匹配。$c_i$和\(\hat{p}_{\hat{\sigma}(i)}(c_i)\)分别表示目标类别标签和预测标签,$b_i$和\(\hat{b}_{\hat{\sigma}}(i)\)分别表示实际边界框参数和预测边界框。在实践中对于$c_i=\Phi$损失缩小$10$倍,以减少类别不平衡。

(2) 模型结构

DETR的结构主要有三部分:

⚪ Backbone

输入图像的尺寸为\(x_{\text{img}} \in \Bbb{R}^{3\times H_0 \times W_0}\),backbone卷积网络生成低分辨率的激活图\(f \in \Bbb{R}^{C\times H \times W}\),通常取\(C=2048,H=\frac{H_0}{32},W=\frac{W_0}{32}\)。

⚪ Transformer encoder

首先使用$1×1$卷积降低通道数为$d$,构造新的特征图\(z_0 \in \Bbb{R}^{d\times H \times W}\)。由于编码器需要序列作为输入,因此将特征图调整为尺寸为$[d,HW]$的序列。由于Transformer置换不变(permutation-invariant)的,在每个注意力层中加入位置编码。DETR编码器从整个图像中提取全局上下文信息。

⚪ Transformer decoder

与标准的Transformer解码器使用自回归模型预测输出序列不同,DETR的解码器并行地将维度为$d$的$N$个输入嵌入解码为$N$个对象的输出嵌入。由于解码器也是置换不变的,因此$N$个输入嵌入必须不同才能产生不同的结果,将其设置为可学习的位置编码,称为对象查询(object query),并加入到每个注意力层的输入中。DETR解码器利用对象之间的成对关系对所有对象进行全局推理。

⚪ Prediction feed-forward networks

前馈网络将$N$个对象的输出嵌入独立解码为边界框的坐标和类别标签,从而得到$N$个最终预测结果。网络包括一个使用ReLU激活函数的$3$层感知机和一个线性映射层,感知机将$N$个预测结果表示为目标边界框的归一化中心坐标和高度、宽度;线性映射层使用softmax函数预测类别标签。由于$N$通常比感兴趣的对象数量大得多,因此预测类别额外增加一个$\Phi$表示未检测到任何对象,类似于目标检测中的“背景”类。

⚪ Auxiliary decoding losses

作者发现训练过程中为解码器增加辅助损失(auxiliary loss)有助于模型输出每个类别的正确目标数量。具体地,在每个解码层后添加参数共享的预测FFN和匈牙利损失,并使用额外的共享layer-norm规范来自不同解码层的预测FFN的输入。

2. 实验分析

作者将DETRFaster R-CNNCOCO数据集上进行性能比较,并对DETR的结构进行详细的消融实验,并将DETR扩展到全景分割任务中。

目标检测和分割实验都是在COCO数据集上进行的,该数据集包含$118$k训练图像和$5$k验证图像。每个图像平均有$7$个实例,单个图像中最多有$63$个实例。目标检测的评估指标采用bbox AP,与Faster R-CNN对比时使用最后一次训练过程中的验证精度,消融实验采用过去$10$轮验证精度的中位数。

DETR使用AdamW算法更新参数,模型中卷积backbone的初始学习率为$10^{-5}$,Transformer的初始学习率为$10^{-4}$,权重衰减为$10^{-4}$,dropout率为$0.1$。Transformer采用Xavier初始化,卷积backbone采用在ImageNet预训练的ResNet模型。作者设计了四种不同的模型:

数据增强使用尺度变换调整图像大小,使短边至少为$480$像素,最多为$800$像素,而长边最多为$1333$像素。为了通过编码器的自注意力机制学习全局关系,在训练期间应用随机裁剪增强,即将图像以$0.5$的概率裁剪为随机矩形子图,然后再次调整为$800-1333$。

在消融实验中,总训练轮数为$300$轮,$200$轮后学习率衰减$10$倍。在性能对比实验中,总训练轮数为$500$轮,$400$轮后学习率衰减$10$倍。

(1) 目标检测实验

作者给出了Faster R-CNN、改进的Faster R-CNN(增加GIoU损失,训练时随机裁剪、$9$倍训练时间)以及DETR的目标检测结果。结果表明,在参数量相当的情况下(41.3M参数中卷积提供23.5M自注意力提供17.8M),DETR检测大目标的性能大幅提高,但检测小目标的性能较差。

(2) 消融实验

⚪ 编码器层数

编码器通过自注意力机制捕捉图像的全局重要性,如果没有编码器,总体AP会下降$3.9$个点。作者认为编码器能够通过全局推理分离每一个目标对象。可视化最后一个编码器层的注意力图,观察到注意力集中在图像中的几个目标上。

⚪ 解码器层数

增加解码器层数会使AP提高。由于DETR预测集合没有冗余,因此不需要NMS。为了验证这一点,在每层解码器后的输出运行一个标准NMS。结果表明NMS提高了来自第一个解码器层的预测性能,这是因为单个解码层无法很好地捕捉输出元素之间的相互关系,因此容易对同一对象进行多个预测。在后续层中NMS带来的改善随着深度的增加而减少。

下图给出了解码器注意力的可视化,注意到解码器的注意力是相当局部的,主要关注物体的四肢,如头部或腿部。这是因为编码器已经通过全局自注意力分离实例,解码器只需要关注局部来提取目标的类别和边界。

⚪ 全连接层

Transformer内部的全连接层可以看作是$1\times 1$卷积层,使得编码器类似于注意力增强卷积网络。将全连接层完全移除后网络参数量从$41.3$M减少到$28.7$M,性能下降了$2.3$ AP,可见全连接层对于取得良好的表现非常重要。

⚪ 位置编码

模型中编码器采用了空间位置编码,解码器采用了空间位置编码和输出位置编码(即目标查询)。作者对固定编码和可学习编码进行消融,并对编码的添加位置进行实验。结果表明,空间位置编码在每一层自注意力层中采用固定编码,输出位置编码在解码器每层自注意力层中采用可学习编码能够取得最好的性能。

⚪ 损失函数

损失函数包括分类损失,l1边界框损失和GIoU边界框损失。作者也给出了不同损失的消融结果:

⚪ 更多分析

作者对解码器中的$20$个查询预测的边界框分布情况进行可视化。每个边界框预测都表示为一个点,其中心坐标位于图像尺寸归一化的方格内。对点进行颜色编码,使绿色对应小方框,红色对应大水平方框,蓝色对应大垂直方框。结果表明每个查询都具有针对不同区域和边界框大小的预测模式,特别地所有查询都能够预测较宽的边界框(绘图中间对齐的红点),这可能与COCO数据集中物体的分布有关。

作者测试了模型的分布外泛化能力。比如训练集中没有出现超过$13$只长颈鹿的图像,于是合成了一个具有$24$个长颈鹿的图像验证模型的泛化能力,结果表明模型能够在图像上找到所有只长颈鹿。

(3) 全景分割实验

DETR是一个通用的可扩展模型,作者在解码器的输出端增加了一个mask headDETR扩展到全景分割任务中。mask head的结构如下:

mask head将每个对象的解码器输出作为输入,计算每个对象的自注意分数,并为每个对象生成具有较小的分辨率的二进制热图。使用FPN结构提升分辨率并获得最终预测结果,并通过DICE/F-1损失和focal损失进行监督学习。

mask head既可以和网络联合训练,也可以单独进行微调训练。实验结果表明两者的表现相似,而微调训练能够缩短总训练时间。为了获得全景分割的最终结果,在每个输出像素的热图上使用argmax确定最终的分割类别,从而确保分割结果没有重叠。

作者给出了不同模型在全景分割任务上的表现,结果表明DETR仍然具有竞争力的表现。