End-to-End Object Detection with Transformers

发布时间 2023-03-26 19:04:33作者: Orange0005

本文提出了一种端到端的,使用transformer的目标检测方法。作者将目标检测视为直接集合预测的问题。相比较于之前的方法,有效地消除了许多手工设计的组件的需求。

之前目标检测中,不论是proposal based的方法,还是anchor based的方法,都需要用到nms(非极大值抑制)等后处理方法来筛选bounding box。但是由于nms的操作,会使得模型调参比较复杂,而且模型部署起来也比较困难。因此,一个端到端的目标检测模型时一直所追求的。

DETR很好的解决了上述问题,不需要proposal和anchors,利用transformer全局建模的能力,把目标检测看成集合预测的问题。并且由于全局建模的能力,DETR不会输出太多冗余的边界框,直接输出最后的bounding box,不需要nms进行后处理,极大的简化了模型。

摘要

将目标检测看成是集合预测任务,不需要nms后处理和生成anchor。DETR提出了两个关键点:一是目标函数,通过二分图匹配的方式,使得模型具有独一无二的预测,也就是没有冗余的框了。二是使用了transformer编码器解码器的架构。

具体的还有两个小细节,一是解码器的输入除了由编码器提供之外,还有一个learned object query,类似于anchors,DETR可以将learned object query和全局图像信息结合起来,通过不停的做注意力操作,从而使得模型直接输出最后的预测框。二是并行的方式,因为transformer2017首先是使用在机器翻译等NLP领域任务中,使用掩码解码器(自回归方式:一个单词一个单词的翻译),而在视觉任务中,并不需要上下文的关系,图像中的目标没有依赖关系,同时也希望越快越好,因此采用了并行的方式,同时输出所有的结果。

DETR的主要优点就是非常的简单,只要支持CNN和Transformer就可以实现,并不需要特殊的库。同时性能也还不错,在COCO数据集上可以和Faster RCNN基线网络打平。此外,DETR可以很简单的扩展到其他任务上,如全景分割。

介绍

目标检测是希望为每个感兴趣的对象预测一组边界框和类别标签。现在的工作都是用一种间接的方式来解决集合预测的任务。如anchors的方法、YOLO的方法,还有物体中心点,FCOS等方法,这些方法会生成冗余的框,就需要使用nms。

首先来看一下DETR的训练过程:

第一步,使用CNN抽取特征。

第二步,使用Transformer编码器去学习全局特征,帮助后面进行检测。

第三步,集合learned oject query用transformer解码器生成N个预测框。

第四步,使用二分匹配,匹配预测框和Ground true(GT)框,在匹配的框里做目标检测的loss。

DETR的推理过程:

第一步,用CNN抽取特征。

第二步,用Transformer编码器去学全局特征,帮助后边做检测。

第三步,结合learned object query用Transformer解码器生成N个预测框。

第四步,置信度大于0.7的作为前景物体保留,其余作为背景。

性能方面,DETR对大物体检测效果比较好,不受限于生成anchor 的大小,这也可能是由于Transformer的全局计算能力所达到的。但DETR在小物体上效果就差一点。另一方面是DETR训练比较慢,作者训练了500个epoch,一般只需十几个epoch。

相关工作

集合预测

集合预测问题的最大困难在于避免重复。大多数现在的检测器都使用后处理(如非极大值抑制)来解决这个问题,但直接集合预测是无后处理的。解决方案是基于匈牙利算法设计一个损失,以在真实值和预测之间创建二分匹配,这样就可以避免重复,建立唯一的对应。并且抛弃了自回归模型,并使用具有并行解码的转换器。

目标检测研究现状

现在研究都是基于初始预测进行检测,two stage 的方法基于proposal,signal stage的方法基于anchors(物体中心点),有证明这些系统的性能在很大程度上取决于这些初始猜测的设置方式。

DETR方法

目标检测的集合预测损失

DETR在单次传递中,可以推断出固定的,N个预测。N可以设置的足够大,明显多于图像中对象的数量,实验中取了100。损失在预测对象和真实对象之间产生最优二分匹配,然后优化对象规格(边界框)损失。匈牙利算法解决二分图匹配问题。Scipy中有linear-sum-assignment函数,输入为cost matrix,输出为最优的方案。

image-20230324095645243

image-20230324095707169

损失包含两部分,分别是类别损失和框的损失。

这里找最优匹配的方式和原来利用先验知识去把预测和proposal和anchors匹配的方式一样,只不过这里的约束更强,一定要得到一个一对一的匹配关系,后续就不需要nms处理。

定义常规的目标检测损失

image-20230324095920748

一旦得到了最佳匹配,即知道生成的100个框中哪个与gt是最优匹配的框,就可以进一步与GT框计算损失函数,然后做梯度回传。

注:作者在这里发现第一部分分类loss去掉log对数,可以使得前后两个损失在大致的取值空间,提升了学习效果。第二部分边界框回归损失,不仅使用了L1-loss(与边界框有关,大框大loss),还使用了generalize iou loss。

DETR模型架构

image-20230324100401312

下面来简单的介绍一下

第一步,输入3×800×1066,经过CNN得到2048×25×34,然后经过1×1的卷积降维得到256×25×34的特征图,加入位置编码后拉长得到850×256(序列长度850,嵌入维度256)。

第二步,进入Transformer编码器得到850×256的输出(这里可以理解为做了全局信息的编码)。

第三步,Transformer的输入为可学习的object queries 100×256(100个框256维度)。输出维度不变100×256。

第四步,通过FFN也就是全连接层,每个框得到6个输出分别对应前景背景概率,框的边界信息(长宽,中心坐标),然后使用匈牙利算法计算最匹配的框,然后根据GT计算梯度,反向回传更新模型。

注:在实验中,作者使用了一些设置,例如,使用固定的位置编码,并将其添加到每一个attention层。并且,对于每一个解码器块的输出,都进行二分匹配,计算最后的损失。

DETR推断

import torch
from torch import nn
from torchvision.models import resnet50

class DETR(nn.Module):
    def __init__(self, num_classes, hidden_dim, nheads,num_encoder_layers, num_decoder_layers):
        super().__init__()
        # We take only convolutional layers from ResNet-50 model
        self.backbone = nn.Sequential(*list(resnet50(pretrained=True).children())[:-2])
        self.conv = nn.Conv2d(2048, hidden_dim, 1)
        self.transformer = nn.Transformer(hidden_dim, nheads,num_encoder_layers, num_decoder_layers)
        self.linear_class = nn.Linear(hidden_dim, num_classes + 1)
        self.linear_bbox = nn.Linear(hidden_dim, 4)
        self.query_pos = nn.Parameter(torch.rand(100, hidden_dim))
        self.row_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
        self.col_embed = nn.Parameter(torch.rand(50, hidden_dim // 2))
        # position embedding
    def forward(self, inputs):
        x = self.backbone(inputs)
        h = self.conv(x)
        H, W = h.shape[-2:]
        pos = torch.cat([self.col_embed[:W].unsqueeze(0).repeat(H, 1, 1), self.row_embed[:H].unsqueeze(1).repeat(1, W, 1),], dim=-1).flatten(0, 1).unsqueeze(1)
        # 位置编码仅添加到了输入中
        h = self.transformer(pos + h.flatten(2).permute(2, 0, 1),self.query_pos.unsqueeze(1))
        return self.linear_class(h), self.linear_bbox(h).sigmoid()

detr = DETR(num_classes=91, hidden_dim=256, nheads=8, num_encoder_layers=6, num_decoder_layers=6
detr.eval()
inputs = torch.randn(1, 3, 800, 1200)
logits, bboxes = detr(inputs)

实验

性能

GFLOPS代表模型大小,FPS代表推理速度。虽然DETR的GFLOPS和参数量相对于Faster RCNN较少,但是推理速度还是稍微慢一点。对于小物体,Faster -RCNN比DETR要高4个点左右,但是对于大物体,DRTE要比Faster-RCNN高出6个点左右。作者认为DETR没有anchors尺寸的限制,并且使用的Transformer具有全局建模能力,所以对大物体比较友好。

可视化

作者将编码器的注意力可视化出来了,每个物体选一个点计算自注意力。我们可以发现,经过Transformer Encoder后每个物体都可以很好的区分开来了,这时候再去做目标检测或者分割任务就简单很多了。


DETR的encoder是学一个全局的特征,让物体之间尽可能分得开。但是对于轮廓点这些细节就需要decoder去做,decoder可以很好的处理遮挡问题。


object query

作者将COCO数据集上得到的所有输出框全都可视化出来。作者将100个object query中20个拿出来,每个正方形代表一个object query。每个object query相当于一个问问题的人,绿色代表小的检测框,红色代表竖向大的检测框,蓝色代表大的横向检测框。例如第一个图,就是不断查询左下角是否有小的目标,中间是否有大的竖向的目标。当经过100个不同object query查询完成后,目标也就检测完成了。

结论

DETR在COCO数据集上与Faster R-CNN基线模型打成平手,并且在分割任务上取得更好的结果。最主要的优势是简单,可以有很大的潜能应用在别的任务上。作者又强调了一下,DETR在大物体上效果非常好。文章存在的缺点作者也自己指出:推理时间有点长、由于使用了Transformer不好优化、小物体上性能也差一些。后续Deformable DETR解决了推理时间和小物体检测差的不足。

论文地址:https://arxiv.org/pdf/2005.12872.pdf