End-to-End Object Detection with Transformers[DETR]


End-to-End Object Detection with Transformers[DETR]

      • 背景
      • 概述
      • 相关技术
        • 输入
        • 提取特征
        • 获取position_embedding
        • transformer
        • encoder
        • decoder
        • 回归
      • 总结

背景

最近在做机器翻译的优化，接触的模型就是transformer， 为了提升性能，在cpu和GPU两个平台c++重新写了整个模型，所以对于机器翻译中transformer的原理细节还是有一定的理解，同时以前做文档图片检索对于图像领域的目标检测也研究颇深，看到最近各大公众号都在推送这篇文章就简单的看了一下，感觉还是蛮有新意的，由于该论文开源，所以直接就跟着代码来解读整篇论文。

概述

在这里插入图片描述
整体来看，该模型首先是经历一个CNN提取特征，然后得到的特征进入transformer, 最后将transformer输出的结果转化为class和box.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
 def forward(self, samples):
        """
        这一段代码时从源码detr.py的DETR中抽出来的代码，为了逻辑清爽，删除了一些
        细枝末节的内容，核心逻辑如下
        """
        #backbone模型中核心就是图中的CNN模型,可以自己选择resnet,vgg什么的，features就是卷积后的输出
        features, pos = self.backbone(samples)#sample 就是图片，大小比如（3,200,250）
        src, mask = features[-1].decompose()
        #transformer模型处理一波
        hs = transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]
        #transformer模型的最终结果为hs，将其分别进入class和box的模型中处理得到class和box
        outputs_class = class_embed(hs)
        outputs_coord = bbox_embed(hs).sigmoid()
        out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}
        return out

下面是大致的推理过程：
在这里插入图片描述

相关技术

输入

作者这里封装了一个类，感觉多此一举，假如我们输入的是如下两张图片，也就说batch为2：
img1 = torch.rand(3, 200, 200),
img2 = torch.rand(3, 200, 250)

1
x = nested_tensor_from_tensor_list([torch.rand(3, 200, 200), torch.rand(3, 200, 250)])

这里会转成nested_tensor, 这个nestd_tensor是什么类型呢？简单说就是把{tensor， mask}打包在一起, tensor就是我么的图片的值，那么mask是什么呢? 当一个batch中的图片大小不一样的时候，我们要把它们处理的整齐，简单说就是把图片都padding成最大的尺寸，padding的方式就是补零，那么batch中的每一张图都有一个mask矩阵，所以mask大小为[2, 200,250], tensor大小为[2,3,200,250]。

提取特征

接下里就是把tensor, 也就是图片输入到特征提取器中，这里作者使用的是残差网络，我做实验的时候用多个resnet-50, 所以tensor经过resnet-50后的结果就是[2,2048,7,8],下面是残差网络最后一层的结构。

(2): Bottleneck(
(conv1): Conv2d(2048, 512, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn1): FrozenBatchNorm2d()
(conv2): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)
(bn2): FrozenBatchNorm2d()
(conv3): Conv2d(512, 2048, kernel_size=(1, 1), stride=(1, 1), bias=False)
(bn3): FrozenBatchNorm2d()
(relu): ReLU(inplace=True)

别忘了，我们还有个mask, mask采用的方式F.interpolate，最后得到的结果是[2,7,8]

获取position_embedding

这里作者使用的三角函数的方式获取position_embediing, 如果你对位置编码不了解，你可以这样理解，“我爱祖国”，“我”位于第一位，如果编码后不加入位置信息，那么“我”这个字的编码信息就是不完善的，所以这里也一样，下面是源码，有兴趣的可以推导一下，position_embediing的输入是上面的NestedTensor={tensor,mask}, 输出最终pos的size为[1,2,256,7,8]。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def forward(self, tensor_list: NestedTensor):
        x = tensor_list.tensors
        mask = tensor_list.mask
        assert mask is not None
        not_mask = ~mask
        y_embed = not_mask.cumsum(1, dtype=torch.float32)
        x_embed = not_mask.cumsum(2, dtype=torch.float32)
        if self.normalize:
            eps = 1e-6
            y_embed = y_embed / (y_embed[:, -1:, :] + eps) * self.scale
            x_embed = x_embed / (x_embed[:, :, -1:] + eps) * self.scale

        dim_t = torch.arange(self.num_pos_feats, dtype=torch.float32, device=x.device)
        dim_t = self.temperature ** (2 * (dim_t // 2) / self.num_pos_feats)

        pos_x = x_embed[:, :, :, None] / dim_t
        pos_y = y_embed[:, :, :, None] / dim_t
        pos_x = torch.stack((pos_x[:, :, :, 0::2].sin(), pos_x[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos_y = torch.stack((pos_y[:, :, :, 0::2].sin(), pos_y[:, :, :, 1::2].cos()), dim=4).flatten(3)
        pos = torch.cat((pos_y, pos_x), dim=3).permute(0, 3, 1, 2)
        return pos

transformer

transformer分为编码和解码，下面分别介绍：

encoder

经过上面一系列操作以后，目前我们拥有src=[ 2, 2048,7,8]，mask=[2,7,8], pos=[1,2,256,7,8]

1
hs = transformer(self.input_proj(src), mask, self.query_embed.weight, pos[-1])[0]#

input_proj：一个卷积层，卷积核为1*1，说白了就是将压缩通道的作用，将2048压缩到256，所以传入transformer的维度是压缩后的[2,256,7,8]。
self.query_embed.weight：现在还用不到，在decoder的时候用的到，到时候再说。
来看一下transformer

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
class Transformer(nn.Module):

    def __init__(self, d_model=512, nhead=8, num_encoder_layers=6,
                 num_decoder_layers=6, dim_feedforward=2048, dropout=0.1,
                 activation="relu", normalize_before=False,
                 return_intermediate_dec=False):
        super().__init__()
        # encode
        # 单层
        encoder_layer = TransformerEncoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        encoder_norm = nn.LayerNorm(d_model) if normalize_before else None
        # 由6个单层组成整个encoder
        self.encoder = TransformerEncoder(encoder_layer, num_encoder_layers, encoder_norm)
        #decode
        decoder_layer = TransformerDecoderLayer(d_model, nhead, dim_feedforward,
                                                dropout, activation, normalize_before)
        decoder_norm = nn.LayerNorm(d_model)
        self.decoder = TransformerDecoder(decoder_layer, num_decoder_layers, decoder_norm,
                                          return_intermediate=return_intermediate_dec)

为了更清楚看到具体模型结构
在这里插入图片描述
根据代码和模型结构可以看到，encoder部分就是6个TransformerEncodeLayer组成，而每一个编码层又由1个self_attention, 2个ffn，2个norm。
在进行encoder之前先还有个处理：

1
2
3
4
bs, c, h, w = src.shape# 这个和我们上面说的一样[2,256,7,8]
src = src.flatten(2).permute(2, 0, 1) # src转为[56,2,256]
pos_embed = pos_embed.flatten(2).permute(2, 0, 1)# pos_embed 转为[56,2,256]
mask = mask.flatten(1) #mask 转为[2,56]

encoder的输入为：src, mask, pos_embed，接下来捋一捋第一个单层encoder的过程

1
2
3
4
5
6
7
8
9
 q = k = self.with_pos_embed(src, pos)# pos + src
 src2 = self.self_attn(q, k, value=src, key_padding_mask=mask)[0]
 #做self_attention,这个不懂的需要补一下transfomer的知识
 src = src + self.dropout1(src2)# 类似于残差网络的加法
 src = self.norm1(src)# norm，这个不是batchnorm，很简单不在详述
 src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))#两个ffn
 src = src + self.dropout2(src2)# 同上残差加法
 src = self.norm2(src)# norm
 return src

根据模型的代码可以看到单层的输出依然为src[56, 2, 256]，第二个单层的输入依然是：src, mask, pos_embed。循环往复6次结束encoder，得到输出memory, memory的size依然为[56, 2, 256].

decoder

encoder结束后我们来看decoder, 先看代码：

1
2
3
tgt = torch.zeros_like(query_embed)
hs = self.decoder(tgt, memory, memory_key_padding_mask=mask,
                  pos=pos_embed, query_pos=query_embed)

现在来找输入：

  1. memory:这个就是encoder的输出，size为[56,2,256]
  2. mask:还是上面的mask
  3. pos_embed:还是上面的pos_embed
  4. query_embed:?
  5. tgt: 每一层的decoder的输入，第一层的话等于0

所以目前我们只要知道query_embed就行了，这个query_embed其实是一个varible，size=[100,2,256],由训练得到，结束后就固定下来了。到目前为止我们获得了decoder的所有输入，和encoder一样我们先来看看单层的decoder的运行流程：

如果你不知道100是啥，那你多少需要看一眼论文，这个100表示将要预测100个目标框，你问为什么是100框，因为作者用的数据集的目标种类有90个，万一一个图上有90个目标你至少都能检测出来吧，所以100个框合理。此外这里和语言模型的输入有很大区别，比如翻译时自回归，也就是说翻译出一个字，然后把这个字作为下一个解码的输入（这里看不懂的可以去看我博客里将transformer的那一篇），作者这里直接用[100, 256]作为输入感觉也是蛮厉害的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
 q = k = self.with_pos_embed(tgt, query_pos)# tgt + query_pos, 第一层的tgt为0
 tgt2 = self.self_attn(q, k, value=tgt, key_padding_mask=mask)[0]# 同上
 tgt = tgt + self.dropout1(tgt2)
 tgt = self.norm1(tgt)
 tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos),
                            key=self.with_pos_embed(memory, pos),
                            value=memory,
                            key_padding_mask=mask)[0]#交叉attention
 tgt = tgt + self.dropout2(tgt2)
 tgt = self.norm2(tgt)
 tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt))))
 tgt = tgt + self.dropout3(tgt2)
 tgt = self.norm3(tgt)
 return tgt

这里的难点可能是交叉attention,也叫encoder_decoder_attention, 这里利用的是encoder的输出来参与计算，里面的计算细节同样可以参考这里，经过上面六次的处理，最后得到的结果为[100,2,256], 返回的时候做一个转换，最终的结果transpose(1, 2)->[100,256,2]。

回归

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
class MLP(nn.Module):
    """ Very simple multi-layer perceptron (also called FFN)"""

    def __init__(self, input_dim, hidden_dim, output_dim, num_layers):
        super().__init__()
        self.num_layers = num_layers
        h = [hidden_dim] * (num_layers - 1)
        self.layers = nn.ModuleList(nn.Linear(n, k) for n, k in zip([input_dim] + h, h + [output_dim]))
    def forward(self, x):
        for i, layer in enumerate(self.layers):
            x = F.relu(layer(x)) if i < self.num_layers - 1 else layer(x)
        return x
       
 self.class_embed = nn.Linear(hidden_dim, num_classes + 1)
 self.bbox_embed = MLP(hidden_dim, hidden_dim, 4, 3)
 
 outputs_class = self.class_embed(hs)
 outputs_coord = self.bbox_embed(hs).sigmoid()
 out = {'pred_logits': outputs_class[-1], 'pred_boxes': outputs_coord[-1]}

这几行代码就不解释了，至于为什么是output_calss[-1], 作为思考题留给大家，如果整个源码撸一遍的话就会知道原因，总的来说最后回归的逻辑比较简单清晰，下面是最后的结果：
pred_logits:[2,100,92]
outputs_coord:[2,100,4]

总结

以上就是整个DETR的推理过程，在训练的时候还涉及到100个框对齐的问题，也不难这里就不再讲述了，如果想彻底理解整个模型，你需要对卷积，attention有比较深刻的理解，不然即使看懂了流程也不明白为什么这样做，该论文的坑位目测还不少，而且对于目标检测的模型来说这个代码量算是少的了，给起来也快，需要毕业的孩纸抓紧啦，哈哈哈