查看原文
其他

搞懂 Vision Transformer 原理和代码,看这篇技术综述就够了

CV开发者都爱看的 极市平台 2022-05-08
↑ 点击蓝字 关注极市平台

作者丨科技猛兽
编辑丨极市平台

极市导读

 

本文对Vision Transformer的原理和代码进行了非常全面详细的解读,一切从Self-attention开始、Transformer的实现和代码以及Transformer+Detection:引入视觉领域的首创DETR。>>加入极市CV技术交流群,走在计算机视觉的最前沿

Transformer 是 Google 的团队在 2017 年提出的一种 NLP 经典模型,现在比较火热的 Bert 也是基于 Transformer。Transformer 模型使用了 Self-Attention 机制,不采用RNN顺序结构,使得模型可以并行化训练,而且能够拥有全局信息。本文将对Vision Transformer的原理和代码进行非常全面的解读。考虑到每篇文章字数的限制,每一篇文章将按照目录的编排包含三个小节,而且这个系列会随着Vision Transformer的发展而长期更新。

目录

(每篇文章对应一个Section,目录持续更新。)

  • Section 1

1 一切从Self-attention开始
1.1 处理Sequence数据的模型
1.2 Self-attention
1.3 Multi-head Self-attention
1.4 Positional Encoding

2 Transformer的实现和代码解读 (NIPS2017)
(来自Google Research, Brain Team)
2.1 Transformer原理分析
2.2 Transformer代码解读

3 Transformer+Detection:引入视觉领域的首创DETR (ECCV2020)
(来自Facebook AI)
3.1 DETR原理分析
3.2 DETR代码解读

  • Section 2

4 Transformer+Detection:Deformable DETR:可变形的Transformer (ICLR2021)
(来自商汤代季峰老师组)
4.1 Deformable DETR原理分析
4.2 Deformable DETR代码解读

5 Transformer+Classification:用于分类任务的Transformer (ICLR2021)
(来自Google Research, Brain Team)
5.1 ViT原理分析
5.2 ViT代码解读

6 Transformer+Image Processing:IPT:用于底层视觉任务的Transformer
(来自北京华为诺亚方舟实验室)
6.1 IPT原理分析

  • Section 3

7 Transformer+Segmentation:SETR:基于Transformer 的语义分割
(来自复旦大学,腾讯优图等)
7.1 SETR原理分析

8 Transformer+GAN:VQGAN:实现高分辨率的图像生成
(来自德国海德堡大学)
8.1 VQGAN原理分析
8.2 VQGAN代码解读

9 Transformer+Distillation:DeiT:高效图像Transformer
(来自Facebook AI)
9.1 DeiT原理分析

1 一切从Self-attention开始

  • 1.1 处理Sequence数据的模型:

Transformer是一个Sequence to Sequence model,特别之处在于它大量用到了self-attention。

要处理一个Sequence,最常想到的就是使用RNN,它的输入是一串vector sequence,输出是另一串vector sequence,如下图1左所示。

如果假设是一个single directional的RNN,那当输出 时,默认 都已经看过了。如果假设是一个bi-directional的RNN,那当输出 时,默认 都已经看过了。RNN非常擅长于处理input是一个sequence的状况。

那RNN有什么样的问题呢?它的问题就在于:RNN很不容易并行化 (hard to parallel)。

为什么说RNN很不容易并行化呢?假设在single directional的RNN的情形下,你今天要算出 ,就必须要先看 再看 再看 再看 ,所以这个过程很难平行处理。

所以今天就有人提出把CNN拿来取代RNN,如下图1右所示。其中,橘色的三角形表示一个filter,每次扫过3个向量 ,扫过一轮以后,就输出了一排结果,使用橘色的小圆点表示。

这是第一个橘色的filter的过程,还有其他的filter,比如图2中的黄色的filter,它经历着与橘色的filter相似的过程,又输出一排结果,使用黄色的小圆点表示。

图1:处理Sequence数据的模型
图2:处理Sequence数据的模型

所以,用CNN,你确实也可以做到跟RNN的输入输出类似的关系,也可以做到输入是一个sequence,输出是另外一个sequence。

但是,表面上CNN和RNN可以做到相同的输入和输出,但是CNN只能考虑非常有限的内容。比如在我们右侧的图中CNN的filter只考虑了3个vector,不像RNN可以考虑之前的所有vector。但是CNN也不是没有办法考虑很长时间的dependency的,你只需要堆叠filter,多堆叠几层,上层的filter就可以考虑比较多的资讯,比如,第二层的filter (蓝色的三角形)看了6个vector,所以,只要叠很多层,就能够看很长时间的资讯。

而CNN的一个好处是:它是可以并行化的 (can parallel),不需要等待红色的filter算完,再算黄色的filter。但是必须要叠很多层filter,才可以看到长时的资讯。所以今天有一个想法:self-attention,如下图3所示,目的是使用self-attention layer取代RNN所做的事情。

图3:You can try to replace any thing that has been done by RNNwith self attention

所以重点是:我们有一种新的layer,叫self-attention,它的输入和输出和RNN是一模一样的,输入一个sequence,输出一个sequence,它的每一个输出 都看过了整个的输入sequence,这一点与bi-directional RNN相同。但是神奇的地方是:它的每一个输出 可以并行化计算。

  • 1.2 Self-attention:

那么self-attention具体是怎么做的呢?

图4:self-attention具体是怎么做的?

首先假设我们的input是图4的 ,是一个sequence,每一个input (vector)先乘上一个矩阵 得到embedding,即向量 。接着这个embedding进入self-attention层,每一个向量 分别乘上3个不同的transformation matrix ,以向量 为例,分别得到3个不同的向量 

图5:self-attention具体是怎么做的?

接下来使用每个query 去对每个key 做attention,attention就是匹配这2个向量有多接近,比如我现在要对 和 做attention,我就可以把这2个向量做scaled inner product,得到 。接下来你再拿 和 做attention,得到 ,你再拿 和 做attention,得到 ,你再拿 和 做attention,得到 。那这个scaled inner product具体是怎么计算的呢?

式中, 是 跟 的维度。因为 的数值会随着dimension的增大而增大,所以要除以 的值,相当于归一化的效果。

接下来要做的事如图6所示,把计算得到的所有 值取 操作。

图6:self-attention具体是怎么做的?

取完 操作以后,我们得到了 ,我们用它和所有的 值进行相乘。具体来讲,把 乘上 ,把 乘上 ,把 乘上 ,把 乘上 ,把结果通通加起来得到 ,所以,今天在产生 的过程中用了整个sequence的资讯 (Considering the whole sequence)。如果要考虑local的information,则只需要学习出相应的 , 就不再带有那个对应分支的信息了;如果要考虑global的information,则只需要学习出相应的 , 就带有全部的对应分支的信息了。

图7:self-attention具体是怎么做的?

同样的方法,也可以计算出 ,如下图8所示, 就是拿query 去对其他的 做attention,得到 ,再与value值 相乘取weighted sum得到的。

图8:self-attention具体是怎么做的?

经过了以上一连串计算,self-attention layer做的事情跟RNN是一样的,只是它可以并行的得到layer输出的结果,如图9所示。现在我们要用矩阵表示上述的计算过程。

图9:self-attention的效果

首先输入的embedding是 ,然后用 乘以transformation matrix 得到 ,它的每一列代表着一个vector 。同理,用 乘以transformation matrix 得到 ,它的每一列代表着一个vector 。用 乘以transformation matrix 得到 ,它的每一列代表着一个vector 

图10:self-attention的矩阵计算过程

接下来是 与 的attention过程,我们可以把vector 横过来变成行向量,与列向量 做内积,这里省略了 。这样, 就成为了 的矩阵,它由4个行向量拼成的矩阵和4个列向量拼成的矩阵做内积得到,如图11所示。

在得到 以后,如上文所述,要得到 , 就要使用 分别与 相乘再求和得到,所以 要再左乘 矩阵。

图11:self-attention的矩阵计算过程

到这里你会发现这个过程可以被表示为,如图12所示:输入矩阵 分别乘上3个不同的矩阵 得到3个中间矩阵 。它们的维度是相同的。把 转置之后与 相乘得到Attention矩阵 ,代表每一个位置两两之间的attention。再将它取 操作得到 ,最后将它乘以 矩阵得到输出vector 

图12:self-attention就是一堆矩阵乘法,可以实现GPU加速
  • 1.3 Multi-head Self-attention:

还有一种multi-head的self-attention,以2个head的情况为例:由 生成的 进一步乘以2个转移矩阵变为 和 ,同理由 生成的 进一步乘以2个转移矩阵变为 和 ,由 生成的 进一步乘以2个转移矩阵变为 和 。接下来 再与 做attention,得到weighted sum的权重 ,再与 做weighted sum得到最终的 。同理得到 。现在我们有了 和 ,可以把它们concat起来,再通过一个transformation matrix调整维度,使之与刚才的 维度一致(这步如图13所示)。

图13:multi-head self-attention
图13:调整b的维度

从下图14可以看到 Multi-Head Attention 包含多个 Self-Attention 层,首先将输入 分别传递到 2个不同的 Self-Attention 中,计算得到 2 个输出结果。得到2个输出矩阵之后,Multi-Head Attention 将它们拼接在一起 (Concat),然后传入一个Linear层,得到 Multi-Head Attention 最终的输出 。可以看到 Multi-Head Attention 输出的矩阵 与其输入的矩阵 的维度是一样的。

图14:multi-head self-attention

这里有一组Multi-head Self-attention的解果,其中绿色部分是一组query和key,红色部分是另外一组query和key,可以发现绿色部分其实更关注global的信息,而红色部分其实更关注local的信息。

图15:Multi-head Self-attention的不同head分别关注了global和local的讯息
  • 1.4 Positional Encoding:

以上是multi-head self-attention的原理,但是还有一个问题是:现在的self-attention中没有位置的信息,一个单词向量的“近在咫尺”位置的单词向量和“远在天涯”位置的单词向量效果是一样的,没有表示位置的信息(No position information in self attention)。所以你输入"A打了B"或者"B打了A"的效果其实是一样的,因为并没有考虑位置的信息。所以在self-attention原来的paper中,作者为了解决这个问题所做的事情是如下图16所示:

图16:self-attention中的位置编码

具体的做法是:给每一个位置规定一个表示位置信息的向量 ,让它与 加在一起之后作为新的 参与后面的运算过程,但是这个向量 是由人工设定的,而不是神经网络学习出来的。每一个位置都有一个不同的 

那到这里一个自然而然的问题是:为什么是 与 相加?为什么不是concatenate?加起来以后,原来表示位置的资讯不就混到 里面去了吗?不就很难被找到了吗?

这里提供一种解答这个问题的思路:

如图15所示,我们先给每一个位置的 append一个one-hot编码的向量 ,得到一个新的输入向量 ,这个向量作为新的输入,乘以一个transformation matrix 。那么:

所以,与 相加就等同于把原来的输入 concat一个表示位置的独热编码 ,再做transformation。

这个与位置编码乘起来的矩阵 是手工设计的,如图17所示。

图17:与位置编码乘起来的转移矩阵WP

Transformer 中除了单词的 Embedding,还需要使用位置 Embedding 表示单词出现在句子中的位置。因为 Transformer 不采用 RNN 的结构,而是使用全局信息,不能利用单词的顺序信息,而这部分信息对于 NLP 来说非常重要。所以 Transformer 中使用位置 Embedding 保存单词在序列中的相对或绝对位置。

位置 Embedding 用 PE表示,PE 的维度与单词 Embedding 是一样的。PE 可以通过训练得到,也可以使用某种公式计算得到。在 Transformer 中采用了后者,计算公式如下:

式中, 表示token在sequence中的位置,例如第一个token "我" 的 

,或者准确意义上是 和 表示了Positional Encoding的维度,的取值范围是: 。所以当 为1时,对应的Positional Encoding可以写成:

式中, 。底数是10000。为什么要使用10000呢,这个就类似于玄学了,原论文中完全没有提啊,这里不得不说说论文的readability的问题,即便是很多高引的文章,最基本的内容都讨论不清楚,所以才出现像上面提问里的讨论,说实话这些论文还远远没有做到easy to follow。这里我给出一个假想:是一个比较接近1的数(1.018),如果用100000,则是1.023。这里只是猜想一下,其实大家应该完全可以使用另一个底数。

这个式子的好处是:

  • 每个位置有一个唯一的positional encoding。
  • 使 能够适应比训练集里面所有句子更长的句子,假设训练集里面最长的句子是有 20 个单词,突然来了一个长度为 21 的句子,则使用公式计算的方法可以计算出第 21 位的 Embedding。
  • 可以让模型容易地计算出相对位置,对于固定长度的间距 ,任意位置的 都可以被 的线性函数表示,因为三角函数特性:

接下来我们看看self-attention在sequence2sequence model里面是怎么使用的,我们可以把Encoder-Decoder中的RNN用self-attention取代掉。

图18:Seq2seq with Attention

2 Transformer的实现和代码解读

  • 2.1 Transformer原理分析:
图19:Transformer

Encoder:

这个图19讲的是一个seq2seq的model,左侧为 Encoder block,右侧为 Decoder block。红色圈中的部分为Multi-Head Attention,是由多个Self-Attention组成的,可以看到 Encoder block 包含一个 Multi-Head Attention,而 Decoder block 包含两个 Multi-Head Attention (其中有一个用到 Masked)。Multi-Head Attention 上方还包括一个 Add & Norm 层,Add 表示残差连接 (Residual Connection) 用于防止网络退化,Norm 表示 Layer Normalization,用于对每一层的激活值进行归一化。比如说在Encoder Input处的输入是机器学习,在Decoder Input处的输入是<BOS>,输出是machine。再下一个时刻在Decoder Input处的输入是machine,输出是learning。不断重复知道输出是句点(.)代表翻译结束。

接下来我们看看这个Encoder和Decoder里面分别都做了什么事情,先看左半部分的Encoder:首先输入 通过一个Input Embedding的转移矩阵 变为了一个张量,即上文所述的 ,再加上一个表示位置的Positional Encoding ,得到一个张量,去往后面的操作。

它进入了这个绿色的block,这个绿色的block会重复 次。这个绿色的block里面有什么呢?它的第1层是一个上文讲的multi-head的attention。你现在一个sequence,经过一个multi-head的attention,你会得到另外一个sequence 

下一个Layer是Add & Norm,这个意思是说:把multi-head的attention的layer的输入 和输出 进行相加以后,再做Layer Normalization,至于Layer Normalization和我们熟悉的Batch Normalization的区别是什么,请参考图20和21。

图20:不同Normalization方法的对比

其中,Batch Normalization和Layer Normalization的对比可以概括为图20,Batch Normalization强行让一个batch的数据的某个channel的 ,而Layer Normalization让一个数据的所有channel的 

图21:Batch Normalization和Layer Normalization的对比

接着是一个Feed Forward的前馈网络和一个Add & Norm Layer。

所以,这一个绿色的block的前2个Layer操作的表达式为:

这一个绿色的block的后2个Layer操作的表达式为:

所以Transformer的Encoder的整体操作为:

Decoder:

现在来看Decoder的部分,输入包括2部分,下方是前一个time step的输出的embedding,即上文所述的 ,再加上一个表示位置的Positional Encoding ,得到一个张量,去往后面的操作。它进入了这个绿色的block,这个绿色的block会重复 次。这个绿色的block里面有什么呢?

首先是Masked Multi-Head Self-attention,masked的意思是使attention只会attend on已经产生的sequence,这个很合理,因为还没有产生出来的东西不存在,就无法做attention。

输出是: 对应 位置的输出词的概率分布。

输入是: 的输出对应 位置decoder的输出。所以中间的attention不是self-attention,它的Key和Value来自encoder,Query来自上一位置 的输出。

解码:这里要特别注意一下,编码可以并行计算,一次性全部Encoding出来,但解码不是一次把所有序列解出来的,而是像 一样一个一个解出来的,因为要用上一个位置的输入当作attention的query。

明确了解码过程之后最上面的图就很好懂了,这里主要的不同就是新加的另外要说一下新加的attention多加了一个mask,因为训练时的output都是Ground Truth,这样可以确保预测第 个位置时不会接触到未来的信息。

  • 包含两个 Multi-Head Attention 层。
  • 第一个 Multi-Head Attention 层采用了 Masked 操作。
  • 第二个 Multi-Head Attention 层的Key,Value矩阵使用 Encoder 的编码信息矩阵 进行计算,而Query使用上一个 Decoder block 的输出计算。
  • 最后有一个 Softmax 层计算下一个翻译单词的概率。

下面详细介绍下Masked Multi-Head Self-attention的具体操作,Masked在Scale操作之后,softmax操作之前

图22:Masked在Scale操作之后,softmax操作之前

因为在翻译的过程中是顺序翻译的,即翻译完第 个单词,才可以翻译第 个单词。通过 Masked 操作可以防止第 个单词知道第 个单词之后的信息。下面以 "我有一只猫" 翻译成 "I have a cat" 为例,了解一下 Masked 操作。在 Decoder 的时候,是需要根据之前的翻译,求解当前最有可能的翻译,如下图所示。首先根据输入 "<Begin>" 预测出第一个单词为 "I",然后根据输入 "<Begin> I" 预测下一个单词 "have"。

Decoder 可以在训练的过程中使用 Teacher Forcing 并且并行化训练,即将正确的单词序列 (<Begin> I have a cat) 和对应输出 (I have a cat <end>) 传递到 Decoder。那么在预测第 个输出时,就要将第  之后的单词掩盖住,注意 Mask 操作是在 Self-Attention 的 Softmax 之前使用的,下面用 0 1 2 3 4 5 分别表示 "<Begin> I have a cat <end>"。

图23:Decoder过程

注意这里transformer模型训练和测试的方法不同:

测试时:

  1. 输入<Begin>,解码器输出 I 。
  2. 输入前面已经解码的<Begin>和 I,解码器输出have。
  3. 输入已经解码的<Begin>,I, have, a, cat,解码器输出解码结束标志位<end>,每次解码都会利用前面已经解码输出的所有单词嵌入信息。

Transformer测试时的解码过程:

训练时:

不采用上述类似RNN的方法 一个一个目标单词嵌入向量顺序输入训练,想采用类似编码器中的矩阵并行算法,一步就把所有目标单词预测出来。要实现这个功能就可以参考编码器的操作,把目标单词嵌入向量组成矩阵一次输入即可。即:并行化训练。

但是在解码have时候,不能利用到后面单词a和cat的目标单词嵌入向量信息,否则这就是作弊(测试时候不可能能未卜先知)。为此引入mask。具体是:在解码器中,self-attention层只被允许处理输出序列中更靠前的那些位置,在softmax步骤前,它会把后面的位置给隐去。

Masked Multi-Head Self-attention的具体操作 如图24所示。

Step1: 输入矩阵包含 "<Begin> I have a cat" (0, 1, 2, 3, 4) 五个单词的表示向量,Mask是一个 5×5 的矩阵。在Mask可以发现单词 0 只能使用单词 0 的信息,而单词 1 可以使用单词 0, 1 的信息,即只能使用之前的信息。输入矩阵 经过transformation matrix变为3个矩阵:Query,Key 和Value 

Step2: 得到 Attention矩阵 ,此时先不急于做softmax的操作,而是先于一个 矩阵相乘,使得attention矩阵的有些位置 归0,得到Masked Attention矩阵 。 矩阵是个下三角矩阵,为什么这样设计?是因为想在计算 矩阵的某一行时,只考虑它前面token的作用。即:在计算 的第一行时,刻意地把 矩阵第一行的后面几个元素屏蔽掉,只考虑 。在产生have这个单词时,只考虑 I,不考虑之后的have a cat,即只会attend on已经产生的sequence,这个很合理,因为还没有产生出来的东西不存在,就无法做attention。

Step3: Masked Attention矩阵进行 Softmax,每一行的和都为 1。但是单词 0 在单词 1, 2, 3, 4 上的 attention score 都为 0。得到的结果再与 矩阵相乘得到最终的self-attention层的输出结果 

Step4: 只是某一个head的结果,将多个head的结果concat在一起之后再最后进行Linear Transformation得到最终的Masked Multi-Head Self-attention的输出结果 

图24:Masked Multi-Head Self-attention的具体操作

第1个Masked Multi-Head Self-attention的 均来自Output Embedding。

第2个Multi-Head Self-attention的 来自第1个Self-attention layer的输出, 来自Encoder的输出。

为什么这么设计? 这里提供一种个人的理解:

来自Transformer Encoder的输出,所以可以看做句子(Sequence)/图片(image)的内容信息(content,比如句意是:"我有一只猫",图片内容是:"有几辆车,几个人等等")

表达了一种诉求:希望得到什么,可以看做引导信息(guide)

通过Multi-Head Self-attention结合在一起的过程就相当于是把我们需要的内容信息指导表达出来

Decoder的最后是Softmax 预测输出单词。因为 Mask 的存在,使得单词 0 的输出 只包含单词 0 的信息。Softmax 根据输出矩阵的每一行预测下一个单词,如下图25所示。

图25:Softmax 根据输出矩阵的每一行预测下一个单词

如下图26所示为Transformer的整体结构。

图26:Transformer的整体结构
  • 2.2 Transformer代码解读:

代码来自:

https://github.com/jadore801120/attention-is-all-you-need-pytorch

ScaledDotProductAttention:
实现的是图22的操作,先令 ,再对结果按位乘以 矩阵,再做 操作,最后的结果与 相乘,得到self-attention的输出。

class ScaledDotProductAttention(nn.Module): ''' Scaled Dot-Product Attention '''
def __init__(self, temperature, attn_dropout=0.1): super().__init__() self.temperature = temperature self.dropout = nn.Dropout(attn_dropout)
def forward(self, q, k, v, mask=None):
attn = torch.matmul(q / self.temperature, k.transpose(2, 3))
if mask is not None: attn = attn.masked_fill(mask == 0, -1e9)
attn = self.dropout(F.softmax(attn, dim=-1)) output = torch.matmul(attn, v)
return output, attn

位置编码 PositionalEncoding:
实现的是式(5)的位置编码。

class PositionalEncoding(nn.Module):
def __init__(self, d_hid, n_position=200): super(PositionalEncoding, self).__init__()
# Not a parameter self.register_buffer('pos_table', self._get_sinusoid_encoding_table(n_position, d_hid))
def _get_sinusoid_encoding_table(self, n_position, d_hid): ''' Sinusoid position encoding table ''' # TODO: make it with torch instead of numpy
def get_position_angle_vec(position): return [position / np.power(10000, 2 * (hid_j // 2) / d_hid) for hid_j in range(d_hid)]
sinusoid_table = np.array([get_position_angle_vec(pos_i) for pos_i in range(n_position)]) sinusoid_table[:, 0::2] = np.sin(sinusoid_table[:, 0::2]) # dim 2i sinusoid_table[:, 1::2] = np.cos(sinusoid_table[:, 1::2]) # dim 2i+1
return torch.FloatTensor(sinusoid_table).unsqueeze(0)#(1,N,d)
def forward(self, x): # x(B,N,d) return x + self.pos_table[:, :x.size(1)].clone().detach()

MultiHeadAttention:
实现图13,14的多头self-attention。

class MultiHeadAttention(nn.Module): ''' Multi-Head Attention module '''
def __init__(self, n_head, d_model, d_k, d_v, dropout=0.1): super().__init__()
self.n_head = n_head self.d_k = d_k self.d_v = d_v
self.w_qs = nn.Linear(d_model, n_head * d_k, bias=False) self.w_ks = nn.Linear(d_model, n_head * d_k, bias=False) self.w_vs = nn.Linear(d_model, n_head * d_v, bias=False) self.fc = nn.Linear(n_head * d_v, d_model, bias=False)
self.attention = ScaledDotProductAttention(temperature=d_k ** 0.5)
self.dropout = nn.Dropout(dropout) self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)

def forward(self, q, k, v, mask=None):
d_k, d_v, n_head = self.d_k, self.d_v, self.n_head sz_b, len_q, len_k, len_v = q.size(0), q.size(1), k.size(1), v.size(1)
residual = q

# Pass through the pre-attention projection: b x lq x (n*dv) # Separate different heads: b x lq x n x dv q = self.w_qs(q).view(sz_b, len_q, n_head, d_k) k = self.w_ks(k).view(sz_b, len_k, n_head, d_k) v = self.w_vs(v).view(sz_b, len_v, n_head, d_v)
# Transpose for attention dot product: b x n x lq x dv q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
if mask is not None: mask = mask.unsqueeze(1) # For head axis broadcasting.
q, attn = self.attention(q, k, v, mask=mask)
#q (sz_b,n_head,N=len_q,d_k) #k (sz_b,n_head,N=len_k,d_k) #v (sz_b,n_head,N=len_v,d_v)
# Transpose to move the head dimension back: b x lq x n x dv # Combine the last two dimensions to concatenate all the heads together: b x lq x (n*dv) q = q.transpose(1, 2).contiguous().view(sz_b, len_q, -1)
#q (sz_b,len_q,n_head,N * d_k) q = self.dropout(self.fc(q)) q += residual
q = self.layer_norm(q)
return q, attn

前向传播Feed Forward Network:

class PositionwiseFeedForward(nn.Module): ''' A two-feed-forward-layer module '''
def __init__(self, d_in, d_hid, dropout=0.1): super().__init__() self.w_1 = nn.Linear(d_in, d_hid) # position-wise self.w_2 = nn.Linear(d_hid, d_in) # position-wise self.layer_norm = nn.LayerNorm(d_in, eps=1e-6) self.dropout = nn.Dropout(dropout)
def forward(self, x):
residual = x
x = self.w_2(F.relu(self.w_1(x))) x = self.dropout(x) x += residual
x = self.layer_norm(x)
return x

EncoderLayer:
实现图26中的一个EncoderLayer,具体的结构如图19所示。

class EncoderLayer(nn.Module): ''' Compose with two layers '''
def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1): super(EncoderLayer, self).__init__() self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout) self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
def forward(self, enc_input, slf_attn_mask=None): enc_output, enc_slf_attn = self.slf_attn( enc_input, enc_input, enc_input, mask=slf_attn_mask) enc_output = self.pos_ffn(enc_output) return enc_output, enc_slf_attn

DecoderLayer:
实现图26中的一个DecoderLayer,具体的结构如图19所示。

class DecoderLayer(nn.Module): ''' Compose with three layers '''
def __init__(self, d_model, d_inner, n_head, d_k, d_v, dropout=0.1): super(DecoderLayer, self).__init__() self.slf_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout) self.enc_attn = MultiHeadAttention(n_head, d_model, d_k, d_v, dropout=dropout) self.pos_ffn = PositionwiseFeedForward(d_model, d_inner, dropout=dropout)
def forward( self, dec_input, enc_output, slf_attn_mask=None, dec_enc_attn_mask=None): dec_output, dec_slf_attn = self.slf_attn( dec_input, dec_input, dec_input, mask=slf_attn_mask) dec_output, dec_enc_attn = self.enc_attn( dec_output, enc_output, enc_output, mask=dec_enc_attn_mask) dec_output = self.pos_ffn(dec_output) return dec_output, dec_slf_attn, dec_enc_attn

Encoder:
实现图26,19左侧的Encoder:

class Encoder(nn.Module): ''' A encoder model with self attention mechanism. '''
def __init__( self, n_src_vocab, d_word_vec, n_layers, n_head, d_k, d_v, d_model, d_inner, pad_idx, dropout=0.1, n_position=200):
super().__init__()
self.src_word_emb = nn.Embedding(n_src_vocab, d_word_vec, padding_idx=pad_idx) self.position_enc = PositionalEncoding(d_word_vec, n_position=n_position) self.dropout = nn.Dropout(p=dropout) self.layer_stack = nn.ModuleList([ EncoderLayer(d_model, d_inner, n_head, d_k, d_v, dropout=dropout) for _ in range(n_layers)]) self.layer_norm = nn.LayerNorm(d_model, eps=1e-6)
def forward(self, src_seq, src_mask, return_attns=False):
enc_slf_attn_list = []
# -- Forward
enc_output = self.dropout(self.position_enc(self.src_word_emb(src_seq))) enc_output = self.layer_norm(enc_output)
for enc_layer in self.layer_stack: enc_output, enc_slf_attn = enc_layer(enc_output, slf_attn_mask=src_mask) enc_slf_attn_list += [enc_slf_attn] if return_attns else []
if return_attns: return enc_output, enc_slf_attn_list return enc_output,

Decoder:
实现图26,19右侧的Decoder:

class Decoder(nn.Module): ''' A decoder model with self attention mechanism. '''

def forward(self, trg_seq, trg_mask, enc_output, src_mask, return_attns=False):
dec_slf_attn_list, dec_enc_attn_list = [], []
# -- Forward dec_output = self.dropout(self.position_enc(self.trg_word_emb(trg_seq))) dec_output = self.layer_norm(dec_output)
for dec_layer in self.layer_stack: dec_output, dec_slf_attn, dec_enc_attn = dec_layer( dec_output, enc_output, slf_attn_mask=trg_mask, dec_enc_attn_mask=src_mask) dec_slf_attn_list += [dec_slf_attn] if return_attns else [] dec_enc_attn_list += [dec_enc_attn] if return_attns else []
if return_attns: return dec_output, dec_slf_attn_list, dec_enc_attn_list return dec_output,

整体结构:
实现图26,19整体的Transformer:

class Transformer(nn.Module): ''' A sequence to sequence model with attention mechanism. '''
def __init__( self, n_src_vocab, n_trg_vocab, src_pad_idx, trg_pad_idx, d_word_vec=512, d_model=512, d_inner=2048, n_layers=6, n_head=8, d_k=64, d_v=64, dropout=0.1, n_position=200, trg_emb_prj_weight_sharing=True, emb_src_trg_weight_sharing=True):
super().__init__()
self.src_pad_idx, self.trg_pad_idx = src_pad_idx, trg_pad_idx
self.encoder = Encoder( n_src_vocab=n_src_vocab, n_position=n_position, d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner, n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v, pad_idx=src_pad_idx, dropout=dropout)
self.decoder = Decoder( n_trg_vocab=n_trg_vocab, n_position=n_position, d_word_vec=d_word_vec, d_model=d_model, d_inner=d_inner, n_layers=n_layers, n_head=n_head, d_k=d_k, d_v=d_v, pad_idx=trg_pad_idx, dropout=dropout)
self.trg_word_prj = nn.Linear(d_model, n_trg_vocab, bias=False)
for p in self.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p)
assert d_model == d_word_vec, \ 'To facilitate the residual connections, \ the dimensions of all module outputs shall be the same.'
self.x_logit_scale = 1. if trg_emb_prj_weight_sharing: # Share the weight between target word embedding & last dense layer self.trg_word_prj.weight = self.decoder.trg_word_emb.weight self.x_logit_scale = (d_model ** -0.5)
if emb_src_trg_weight_sharing: self.encoder.src_word_emb.weight = self.decoder.trg_word_emb.weight

def forward(self, src_seq, trg_seq):
src_mask = get_pad_mask(src_seq, self.src_pad_idx) trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)
enc_output, *_ = self.encoder(src_seq, src_mask) dec_output, *_ = self.decoder(trg_seq, trg_mask, enc_output, src_mask) seq_logit = self.trg_word_prj(dec_output) * self.x_logit_scale
return seq_logit.view(-1, seq_logit.size(2))

产生Mask:

def get_pad_mask(seq, pad_idx): return (seq != pad_idx).unsqueeze(-2)

def get_subsequent_mask(seq): ''' For masking out the subsequent info. ''' sz_b, len_s = seq.size() subsequent_mask = (1 - torch.triu( torch.ones((1, len_s, len_s), device=seq.device), diagonal=1)).bool() return subsequent_mask

src_mask = get_pad_mask(src_seq, self.src_pad_idx)
用于产生Encoder的Mask,它是一列Bool值,负责把标点mask掉。
trg_mask = get_pad_mask(trg_seq, self.trg_pad_idx) & get_subsequent_mask(trg_seq)
用于产生Decoder的Mask。它是一个矩阵,如图24中的Mask所示,功能已在上文介绍。

3 Transformer+Detection:引入视觉领域的首创DETR

论文名称:End-to-End Object Detection with Transformers

论文地址:

https://arxiv.org/abs/2005.12872arxiv.org

  • 3.1 DETR原理分析:

本文的任务是Object detection,用到的工具是Transformers,特点是End-to-end。

目标检测的任务是要去预测一系列的Bounding Box的坐标以及Label, 现代大多数检测器通过定义一些proposal,anchor或者windows,把问题构建成为一个分类和回归问题来间接地完成这个任务。文章所做的工作,就是将transformers运用到了object detection领域,取代了现在的模型需要手工设计的工作,并且取得了不错的结果。在object detection上DETR准确率和运行时间上和Faster RCNN相当;将模型 generalize 到 panoptic segmentation 任务上,DETR表现甚至还超过了其他的baseline。DETR第一个使用End to End的方式解决检测问题,解决的方法是把检测问题视作是一个set prediction problem,如下图27所示。

图27:DETR结合CNN和Transformer的结构,并行实现预测

网络的主要组成是CNN和Transformer,Transformer借助第1节讲到的self-attention机制,可以显式地对一个序列中的所有elements两两之间的interactions进行建模,使得这类transformer的结构非常适合带约束的set prediction的问题。DETR的特点是:一次预测,端到端训练,set loss function和二分匹配。

文章的主要有两个关键的部分。

第一个是用transformer的encoder-decoder架构一次性生成 个box prediction。其中 是一个事先设定的、比远远大于image中object个数的一个整数。

第二个是设计了bipartite matching loss,基于预测的boxex和ground truth boxes的二分图匹配计算loss的大小,从而使得预测的box的位置和类别更接近于ground truth。

DETR整体结构可以分为四个部分:backbone,encoder,decoder和FFN,如下图28所示,以下分别解释这四个部分:

图28:DETR整体结构

1 首先看backbone: CNN backbone处理 维的图像,把它转换为维的feature map(一般来说 ),backbone只做这一件事。

2 再看encoder: encoder的输入是维的feature map,接下来依次进行以下过程:

  • 通道数压缩: 先用 convolution处理,将channels数量从 压缩到 ,即得到维的新feature map。
  • 转化为序列化数据: 将空间的维度(高和宽)压缩为一个维度,即把上一步得到的维的feature map通过reshape成维的feature map。
  • 位置编码: 在得到了维的feature map之后,正式输入encoder之前,需要进行 Positional Encoding 。这一步在第2节讲解transformer的时候已经提到过,因为在self-attention中需要有表示位置的信息,否则你的sequence = "A打了B" 还是sequence = "B打了A"的效果是一样的。但是transformer encoder这个结构本身却无法体现出位置信息。也就是说,我们需要对这个 维的feature map做positional encoding。

进行完位置编码以后根据paper中的图片会有个相加的过程,如下图问号处所示。很多读者有疑问的地方是:论文图示中相加的2个张量,一个是input embedding,另一个是位置编码维度看上去不一致,是怎么相加的?后面会解答。

图:怎么相加的?

原版Transformer和Vision Transformer (第4节讲述)的Positional Encoding的表达式为:

式中, 就是这个 维的feature map的第一维, 。表示token在sequence中的位置,sequence的长度是 ,例如第一个token 的 

,或者准确意义上是 和 表示了Positional Encoding的维度,的取值范围是: 。所以当 为1时,对应的Positional Encoding可以写成:

式中, 

第一点不同的是,原版Transformer只考虑 方向的位置编码,但是DETR考虑了 方向的位置编码,因为图像特征是2-D特征。采用的依然是 模式,但是需要考虑 两个方向。不是类似vision transoformer做法简单的将其拉伸为 ,然后从 进行长度为256的位置编码,而是考虑了 方向同时编码,每个方向各编码128维向量,这种编码方式更符合图像特点。

Positional Encoding的输出张量是: ,其中 代表位置编码的长度, 代表张量的位置。意思是说,这个特征图上的任意一个点 有个位置编码,这个编码的长度是256,其中,前128维代表 的位置编码,后128维代表 的位置编码。

假设你想计算任意一个位置 的Positional Encoding,把 代入(11)式的 式和 式可以计算得到128维的向量,它代表 的位置编码,再把 代入(11)式的 式和 式可以计算得到128维的向量,它代表 的位置编码,把这2个128维的向量拼接起来,就得到了一个256维的向量,它代表 的位置编码。

计算所有位置的编码,就得到了 的张量,代表这个batch的位置编码。编码矩阵的维度是 ,也把它序列化成维度为 维的张量。

准备与维的feature map相加以后输入Encoder。

值得注意的是,网上许多解读文章没有搞清楚 "转化为序列化数据"这一步和 "位置编码"的顺序关系,以及变量的shape到底是怎样变化的,这里我用一个图表达,终结这个问题。

图29:变量的shape的变化,变量一律使用方块表达。

所以,了解了DETR的位置编码之后,你应该明白了其实input embedding和位置编码维度其实是一样的,只是论文图示为了突出二位编码所以画的不一样罢了,如下图所示:

图:input embedding与positional embedding的shape是一致的

另一点不同的是,原版Transformer 只在Encoder之前使用了Positional Encoding,而且是在输入上进行Positional Encoding,再把输入经过transformation matrix变为Query,Key和Value这几个张量。但是DETR在Encoder的每一个Multi-head Self-attention之前都使用了Positional Encoding,且只对Query和Key使用了Positional Encoding,即:只把维度为维的位置编码与维度为维的Query和Key相加,而不与Value相加。

如图30所示为DETR的Transformer的详细结构,读者可以对比下原版Transformer的结构,如图19所示,为了阅读的方便我把图19又贴在下面了。

可以发现,除了Positional Encoding设置的不一样外,Encoder其他的结构是一致的。每个Encoder Layer包含一个multi-head self-attention 的module和一个前馈网络Feed Forward Network。

Encoder最终输出的是 维的编码矩阵Embedding,按照原版Transformer的做法,把这个东西给Decoder。

总结下和原始transformer编码器不同的地方:

  • 输入编码器的位置编码需要考虑2-D空间位置。
  • 位置编码向量需要加入到每个Encoder Layer中。
  • 在编码器内部位置编码Positional Encoding仅仅作用于Query和Key,即只与Query和Key相加,Value不做任何处理。
图30:Transformer详细结构。为了方便理解,我把每个变量的维度标在了图上。
图19:Transformer整体结构

3 再看decoder:

DETR的Decoder和原版Transformer的decoder是不太一样的,如下图30和19所示。

先回忆下原版Transformer,看下图19的decoder的最后一个框:output probability,代表我们一次只产生一个单词的softmax,根据这个softmax得到这个单词的预测结果。这个过程我们表达为:predicts the output sequence one element at a time

不同的是,DETR的Transformer Decoder是一次性处理全部的object queries,即一次性输出全部的predictions;而不像原始的Transformer是auto-regressive的,从左到右一个词一个词地输出。这个过程我们表达为:decodes the N objects in parallel at each decoder layer。

DETR的Decoder主要有两个输入:

  1. Transformer Encoder输出的Embedding与 position encoding 之和。
  2. Object queries。

其中,Embedding就是上文提到的 的编码矩阵。这里着重讲一下Object queries。

Object queries是一个维度为 维的张量,数值类型是nn.Embedding,说明这个张量是可以学习的,即:我们的Object queries是可学习的。Object queries矩阵内部通过学习建模了100个物体之间的全局关系,例如房间里面的桌子旁边(A类)一般是放椅子(B类),而不会是放一头大象(C类),那么在推理时候就可以利用该全局注意力更好的进行解码预测输出。

Decoder的输入一开始也初始化成维度为 维的全部元素都为0的张量,和Object queries加在一起之后充当第1个multi-head self-attention的Query和Key。第一个multi-head self-attention的Value为Decoder的输入,也就是全0的张量。

到了每个Decoder的第2个multi-head self-attention,它的Key和Value来自Encoder的输出张量,维度为 ,其中Key值还进行位置编码。Query值一部分来自第1个Add and Norm的输出,维度为 的张量,另一部分来自Object queries,充当可学习的位置编码。所以,第2个multi-head self-attention的Key和Value的维度为 ,而Query的维度为

每个Decoder的输出维度为 ,送入后面的前馈网络,具体的变量维度的变化见图30。

到这里你会发现:Object queries充当的其实是位置编码的作用,只不过它是可以学习的位置编码,所以,我们对Encoder和Decoder的每个self-attention的Query和Key的位置编码做个归纳,如图31所示,Value没有位置编码:

图31:Transformer的位置编码来自哪里?

得到了Decoder的输出以后,如前文所述,应该是输出维度为 的张量。接下来要送入2个前馈网络FFN得到class和Bounding Box。它们会得到 个预测目标,包含类别和Bounding Box,当然这个100肯定是大于图中的目标总数的。如果不够100,则采用背景填充,计算loss时候回归分支分支仅仅计算有物体位置,背景集合忽略。所以,DETR输出张量的维度为输出的张量的维度是 和 。对应COCO数据集来说, , 指的是每个预测目标归一化的 。归一化就是除以图片宽高进行归一化。

到这里我们了解了DETR的网络架构,我们发现,它输出的张量的维度是 分类分支:回归分支: ,其中,前者是指100个预测框的类型,后者是指100个预测框的Bounding Box,但是读者可能会有疑问:预测框和真值是怎么一一对应的?换句话说:你怎么知道第47个预测框对应图片里的狗,第88个预测框对应图片里的车?等等。

我们下面就来聊聊这个问题。

相比Faster R-CNN等做法,DETR最大特点是将目标检测问题转化为无序集合预测问题(set prediction)。论文中特意指出Faster R-CNN这种设置一大堆anchor,然后基于anchor进行分类和回归其实属于代理做法即不是最直接做法,目标检测任务就是输出无序集合,而Faster R-CNN等算法通过各种操作,并结合复杂后处理最终才得到无序集合属于绕路了,而DETR就比较纯粹了。现在核心问题来了:输出的 个检测结果是无序的,如何和 计算loss?这就需要用到经典的双边匹配算法了,也就是常说的匈牙利算法,该算法广泛应用于最优分配问题。

一幅图片,我们把第 个物体的真值表达为 ,其中, 表示它的 , 表示它的 。我们定义 为网络输出的 个预测值。

假设我们已经了解了什么是匈牙利算法(先假装了解了),对于第 个 , 为匈牙利算法得到的与 对应的prediction的索引。我举个栗子,比如 ,意思就是:与第3个真值对应的预测值是第18个。

那我能根据 ,找到 那究竟是如何找到呢?

我们看看这个表达式是甚么意思,对于某一个真值 ,假设我们已经找到这个真值对应的预测值 ,这里的 是所有可能的排列,代表从真值索引到预测值索引的所有的映射,然后用 最小化 和 的距离。这个 具体是:

意思是:假设当前从真值索引到预测值索引的所有的映射为 ,对于图片中的每个真值 ,先找到对应的预测值 ,再看看分类网络的结果 ,取反作为 的第1部分。再计算回归网络的结果 与真值的 的差异,即 ,作为 的第2部分。

所以,可以使得 最小的排列 就是我们要找的排列,即:对于图片中的每个真值 来讲, 就是这个真值所对应的预测值的索引。

请读者细品这个 寻找匹配的过程 ,这就是匈牙利算法的过程。是不是与Anchor或Proposal有异曲同工的地方,只是此时我们找的是一对一匹配。

接下来就是使用上一步得到的排列 ,计算匈牙利损失:

式中的 具体为:

最常用的 对于大小 会有不同的标度,即使它们的相对误差是相似的。为了缓解这个问题,作者使用了 和广义IoU损耗 的线性组合,它是比例不变的。

Hungarian意思就是匈牙利,也就是前面的 ,上述意思是需要计算 个 和 个输预测出集合两两之间的广义距离,距离越近表示越可能是最优匹配关系,也就是两者最密切。广义距离的计算考虑了分类分支和回归分支。

最后,再概括一下DETR的End-to-End的原理,前面那么多段话就是为了讲明白这个事情,如果你对前面的论述还存在疑问的话,把下面一直到Experiments之前的这段话看懂就能解决你的困惑。

DETR是怎么训练的?

训练集里面的任何一张图片,假设第1张图片,我们通过模型产生100个预测框 ,假设这张图片有只3个 ,它们分别是 

问题是:我怎么知道这100个预测框哪个是对应 ,哪个是对应 ,哪个是对应 

我们建立一个 的矩阵,矩阵里面的元素就是 式的计算结果,举个例子:比如左上角的 号元素的含义是:第1个预测框对应 的情况下的 值。我们用scipy.optimize 这个库中的 linear_sum_assignment 函数找到最优的匹配,这个过程我们称之为:"匈牙利算法 (Hungarian Algorithm)"

假设linear_sum_assignment 做完以后的结果是:第 个预测框对应 ,第 个预测框对应 ,第 个预测框对应 

现在把第 个预测框挑出来,按照 式计算Loss,得到这个图片的Loss。

把所有的图片按照这个模式去训练模型。

训练完以后怎么用?

训练完以后,你的模型学习到了一种能力,即:模型产生的100个预测框,它知道某个预测框该对应什么 ,比如,模型学习到:第1个 对应 ,第2个 对应 ,第3个 对应 ,第4个 对应 ,第5个 对应 ,第6-100个 对应 ,等等。

以上只是我举的一个例子,意思是说:模型知道了自己的100个预测框每个该做什么事情,即:每个框该预测什么样的 

为什么训练完以后,模型学习到了一种能力,即:模型产生的100个预测框,它知道某个预测框该对应什么 

还记得前面说的Object queries吗?它是一个维度为 维的张量,初始时元素全为 。实现方式是nn.Embedding(num_queries, hidden_dim),这里num_queries=100,hidden_dim=256,它是可训练的。这里的 指的是batch size,我们考虑单张图片,所以假设Object queries是一个维度为 维的张量。我们训练完模型以后,这个张量已经训练完了,那此时的Object queries究竟代表什么?

我们把此时的Object queries看成100个格子,每个格子是个256维的向量。训练完以后,这100个格子里面注入了不同 的位置信息和类别信息。比如第1个格子里面的这个256维的向量代表着 这种 的位置信息,这种信息是通过训练,考虑了所有图片的某个位置附近的 编码特征,属于和位置有关的全局 统计信息。

测试时,假设图片中有 三种物体,该图片会输入到编码器中进行特征编码,假设特征没有丢失,Decoder的KeyValue就是编码器输出的编码向量(如图30所示),而Query就是Object queries,就是我们的100个格子。

Query可以视作代表不同 的信息,而Key和Value可以视作代表图像的全局信息。

现在通过注意力模块将QueryKey计算,然后加权Value得到解码器输出。对于第1个格子的Query会和Key中的所有向量进行计算,目的是查找某个位置附近有没有 ,如果有那么该特征就会加权输出,对于第3个格子的Query会和Key中的所有向量进行计算,目的是查找某个位置附近有没有 ,很遗憾,这个没有,所以输出的信息里面没有 

整个过程计算完成后就可以把编码向量中的 的编码嵌入信息提取出来,然后后面接 进行分类和回归就比较容易,因为特征已经对齐了。

发现了吗?Object queries在训练过程中对于 个格子会压缩入对应的和位置和类别相关的统计信息,在测试阶段就可以利用该Query去和某个图像的编码特征Key,Value计算,若图片中刚好有Query想找的特征,比如 ,则这个特征就能提取出来,最后通过2个 进行分类和回归。所以前面才会说Object queries作用非常类似Faster R-CNN中的anchor,这个anchor是可学习的,由于维度比较高,故可以表征的东西丰富,当然维度越高,训练时长就会越长。

这就是DETR的End-to-End的原理,可以简单归结为上面的几段话,你读懂了上面的话,也就明白了DETR以及End-to-End的Detection模型原理。

Experiments:

1. 性能对比:

图32:DETR与Fast R-CNN的性能对比

2. 编码器层数对比实验:

图33:编码器数目与模型性能

可以发现,编码器层数越多越好,最后就选择6。

下图34为最后一个Encoder Layer的attention可视化,Encoder已经分离了instances,简化了Decoder的对象提取和定位。

图34:最后一个Encoder Layer的attention可视化

3. 解码器层数对比实验:

图35:每个Decoder Layer后的AP和AP 50性能。

可以发现,性能随着解码器层数的增加而提升,DETR本不需要NMS,但是作者也进行了,上图中的NMS操作是指DETR的每个解码层都可以输入无序集合,那么将所有解码器无序集合全部保留,然后进行NMS得到最终输出,可以发现性能稍微有提升,特别是AP50。这可以通过以下事实来解释:Transformer的单个Decoder Layer不能计算输出元素之间的任何互相关,因此它易于对同一对象进行多次预测。在第2个和随后的Decoder Layer中,self-attention允许模型抑制重复预测。所以NMS带来的改善随着Decoder Layer的增加而减少。在最后几层,作者观察到AP的一个小损失,因为NMS错误地删除了真实的positive prediction。

图36:Decoder Layer的attention可视化

类似于可视化编码器注意力,作者在图36中可视化解码器注意力,用不同的颜色给每个预测对象的注意力图着色。观察到,解码器的attention相当局部,这意味着它主要关注对象的四肢,如头部或腿部。我们假设,在编码器通过全局关注分离实例之后,解码器只需要关注极端来提取类和对象边界。

  • 3.2 DETR代码解读:

https://github.com/facebookresearch/detr

分析都注释在了代码中。

二维位置编码:
DETR的二维位置编码:
首先构造位置矩阵x_embed和y_embed,这里用到了python函数cumsum,作用是对一个矩阵的元素进行累加,那么累加以后最后一个元素就是所有累加元素的和,省去了求和的步骤,直接用这个和做归一化,对应x_embed[:, :, -1:]和y_embed[:, -1:, :]。
这里我想着重强调下代码中一些变量的shape,方便读者掌握作者编程的思路:
值得注意的是,tensor_list的类型是NestedTensor,内部自动附加了mask,用于表示动态shape,是pytorch中tensor新特性https://github.com/pytorch/nestedtensor。全是false。
x:(b,c,H,W)
mask:(b,H,W),全是False。
not_mask:(b,H,W),全是True。
首先出现的y_embed:(b,H,W),具体是1,1,1,1,......,2,2,2,2,......3,3,3,3,......
首先出现的x_embed:(b,H,W),具体是1,2,3,4,......,1,2,3,4,......1,2,3,4,......
self.num_pos_feats = 128
首先出现的dim_t = [0,1,2,3,.....,127]
pos_x:(b,H,W,128)
pos_y:(b,H,W,128)
flatten后面的数字指的是:flatten()方法应从哪个轴开始展开操作。
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)
这一步执行完以后变成(b,H,W,2,64)通过flatten()方法从第3个轴开始展平,变为:(b,H,W,128)
torch.cat((pos_y, pos_x), dim=3)之后变为(b,H,W,256),再最后permute为(b,256,H,W)。
PositionEmbeddingSine类继承nn.Module类。

class PositionEmbeddingSine(nn.Module):
def __init__(self, num_pos_feats=64, temperature=10000, normalize=False, scale=None): super().__init__() self.num_pos_feats = num_pos_feats self.temperature = temperature self.normalize = normalize if scale is not None and normalize is False: raise ValueError("normalize should be True if scale is passed") if scale is None: scale = 2 * math.pi self.scale = scale
def forward(self, tensor_list: NestedTensor):#输入是b,c,h,w#tensor_list的类型是NestedTensor,内部自动附加了mask,#用于表示动态shape,是pytorch中tensor新特性https://github.com/pytorch/nestedtensor x = tensor_list.tensors# 附加的mask,shape是b,h,w 全是false mask = tensor_list.mask assert mask is not None not_mask = ~mask# 因为图像是2d的,所以位置编码也分为x,y方向# 1 1 1 1 .. 2 2 2 2... 3 3 3... y_embed = not_mask.cumsum(1, dtype=torch.float32)# 1 2 3 4 ... 1 2 3 4... 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
# num_pos_feats = 128# 0~127 self.num_pos_feats=128,因为前面输入向量是256,编码是一半sin,一半cos 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)
# 输出shape=b,h,w,128 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)# 每个特征图的xy位置都编码成256的向量,其中前128是y方向编码,而128是x方向编码 return pos# b,n=256,h,w

作者定义了一种数据结构:NestedTensor,里面打包存了两个变量:x 和mask。

NestedTensor:
里面打包存了两个变量:x 和mask。
to()函数:把变量移到GPU中。

Backbone:

class BackboneBase(nn.Module):
def __init__(self, backbone: nn.Module, train_backbone: bool, num_channels: int, return_interm_layers: bool): super().__init__() for name, parameter in backbone.named_parameters(): if not train_backbone or 'layer2' not in name and 'layer3' not in name and 'layer4' not in name: parameter.requires_grad_(False) if return_interm_layers: return_layers = {"layer1": "0", "layer2": "1", "layer3": "2", "layer4": "3"} else: return_layers = {'layer4': "0"}
#作用的模型:定义BackboneBase时传入的nn.Moduleclass的backbone,返回的layer:来自bool变量return_interm_layers self.body = IntermediateLayerGetter(backbone, return_layers=return_layers) self.num_channels = num_channels
def forward(self, tensor_list: NestedTensor):#BackboneBase的输入是一个NestedTensor#xs中间层的输出, xs = self.body(tensor_list.tensors) out: Dict[str, NestedTensor] = {} for name, x in xs.items(): m = tensor_list.mask assert m is not None#F.interpolate上下采样,调整mask的size#to(torch.bool) 把mask转化为Bool型变量 mask = F.interpolate(m[None].float(), size=x.shape[-2:]).to(torch.bool)[0] out[name] = NestedTensor(x, mask) return out

class Backbone(BackboneBase): """ResNet backbone with frozen BatchNorm.""" def __init__(self, name: str, train_backbone: bool, return_interm_layers: bool, dilation: bool):#根据name选择backbone, num_channels, return_interm_layers等,传入BackboneBase初始化 backbone = getattr(torchvision.models, name)( replace_stride_with_dilation=[False, False, dilation], pretrained=is_main_process(), norm_layer=FrozenBatchNorm2d) num_channels = 512 if name in ('resnet18', 'resnet34') else 2048 super().__init__(backbone, train_backbone, num_channels, return_interm_layers)

把Backbone和之前的PositionEmbeddingSine连在一起:
Backbone完以后输出(b,c,h,w),再经过PositionEmbeddingSine输出(b,H,W,256)。

class Joiner(nn.Sequential): def __init__(self, backbone, position_embedding): super().__init__(backbone, position_embedding)
def forward(self, tensor_list: NestedTensor): xs = self[0](tensor_list) out: List[NestedTensor] = [] pos = [] for name, x in xs.items(): out.append(x) # position encoding pos.append(self[1](x).to(x.tensors.dtype))
return out, pos

def build_backbone(args):#position_embedding是个nn.module position_embedding = build_position_encoding(args) train_backbone = args.lr_backbone > 0 return_interm_layers = args.masks#backbone是个nn.module backbone = Backbone(args.backbone, train_backbone, return_interm_layers, args.dilation)#nn.Sequential在一起 model = Joiner(backbone, position_embedding) model.num_channels = backbone.num_channels return model

Transformer的一个Encoder Layer:

class TransformerEncoderLayer(nn.Module):
def __init__(self, d_model, nhead, dim_feedforward=2048, dropout=0.1, activation="relu", normalize_before=False): super().__init__() self.self_attn = nn.MultiheadAttention(d_model, nhead, dropout=dropout) # Implementation of Feedforward model self.linear1 = nn.Linear(d_model, dim_feedforward) self.dropout = nn.Dropout(dropout) self.linear2 = nn.Linear(dim_feedforward, d_model)
self.norm1 = nn.LayerNorm(d_model) self.norm2 = nn.LayerNorm(d_model) self.dropout1 = nn.Dropout(dropout) self.dropout2 = nn.Dropout(dropout)
self.activation = _get_activation_fn(activation) self.normalize_before = normalize_before
def with_pos_embed(self, tensor, pos: Optional[Tensor]): return tensor if pos is None else tensor + pos
def forward_post(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): # 和标准做法有点不一样,src加上位置编码得到q和k,但是v依然还是src, # 也就是v和qk不一样 q = k = self.with_pos_embed(src, pos) src2 = self.self_attn(q, k, value=src, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0]#Add and Norm src = src + self.dropout1(src2) src = self.norm1(src)#FFN src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))#Add and Norm src = src + self.dropout2(src2) src = self.norm2(src) return src
def forward_pre(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): src2 = self.norm1(src) q = k = self.with_pos_embed(src2, pos) src2 = self.self_attn(q, k, value=src2, attn_mask=src_mask, key_padding_mask=src_key_padding_mask)[0] src = src + self.dropout1(src2) src2 = self.norm2(src) src2 = self.linear2(self.dropout(self.activation(self.linear1(src2)))) src = src + self.dropout2(src2) return src
def forward(self, src, src_mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): if self.normalize_before: return self.forward_pre(src, src_mask, src_key_padding_mask, pos) return self.forward_post(src, src_mask, src_key_padding_mask, pos)

有了一个Encoder Layer的定义,再看Transformer的整个Encoder:

class TransformerEncoder(nn.Module): def __init__(self, encoder_layer, num_layers, norm=None): super().__init__() # 编码器copy6份 self.layers = _get_clones(encoder_layer, num_layers) self.num_layers = num_layers self.norm = norm
def forward(self, src, mask: Optional[Tensor] = None, src_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None): # 内部包括6个编码器,顺序运行 # src是图像特征输入,shape=hxw,b,256 output = src for layer in self.layers: # 每个编码器都需要加入pos位置编码 # 第一个编码器输入来自图像特征,后面的编码器输入来自前一个编码器输出 output = layer(output, src_mask=mask, src_key_padding_mask=src_key_padding_mask, pos=pos) return output

Object Queries:可学习的位置编码:
注释中已经注明了变量的shape的变化过程,最终输出的是与Positional Encoding维度相同的位置编码,维度是(b,H,W,256),只是现在这个位置编码是可学习的了。

class PositionEmbeddingLearned(nn.Module): """ Absolute pos embedding, learned. """ def __init__(self, num_pos_feats=256): super().__init__()]#这里使用了nn.Embedding,这是一个矩阵类,里面初始化了一个随机矩阵,矩阵的长是字典的大小,宽是用来表示字典中每个元素的属性向量,# 向量的维度根据你想要表示的元素的复杂度而定。类实例化之后可以根据字典中元素的下标来查找元素对应的向量。输入下标0,输出就是embeds矩阵中第0行。 self.row_embed = nn.Embedding(50, num_pos_feats) self.col_embed = nn.Embedding(50, num_pos_feats) self.reset_parameters()
def reset_parameters(self): nn.init.uniform_(self.row_embed.weight) nn.init.uniform_(self.col_embed.weight)
#输入依旧是NestedTensor def forward(self, tensor_list: NestedTensor): x = tensor_list.tensors h, w = x.shape[-2:] i = torch.arange(w, device=x.device) j = torch.arange(h, device=x.device)
#x_emb:(w, 128)#y_emb:(h, 128) x_emb = self.col_embed(i) y_emb = self.row_embed(j) pos = torch.cat([ x_emb.unsqueeze(0).repeat(h, 1, 1),#(1,w,128) → (h,w,128) y_emb.unsqueeze(1).repeat(1, w, 1),#(h,1,128) → (h,w,128) ], dim=-1).permute(2, 0, 1).unsqueeze(0).repeat(x.shape[0], 1, 1, 1)#(h,w,256) → (256,h,w) → (1,256,h,w) → (b,256,h,w) return pos

def build_position_encoding(args): N_steps = args.hidden_dim // 2 if args.position_embedding in ('v2', 'sine'): # TODO find a better way of exposing other arguments position_embedding = PositionEmbeddingSine(N_steps, normalize=True) elif args.position_embedding in ('v3', 'learned'): position_embedding = PositionEmbeddingLearned(N_steps) else: raise ValueError(f"not supported {args.position_embedding}")
return position_embedding

Transformer的一个Decoder Layer:
注意变量的命名:
object queries(query_pos)
Encoder的位置编码(pos)
Encoder的输出(memory)

def forward_post(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None):
#query,key的输入是object queries(query_pos) + Decoder的输入(tgt),shape都是(100,b,256)#value的输入是Decoder的输入(tgt),shape = (100,b,256) q = k = self.with_pos_embed(tgt, query_pos)
#Multi-head self-attention tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0]#Add and Norm tgt = tgt + self.dropout1(tgt2) tgt = self.norm1(tgt)
#query的输入是上一个attention的输出(tgt) + object queries(query_pos)#key的输入是Encoder的位置编码(pos) + Encoder的输出(memory)#value的输入是Encoder的输出(memory) tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt, query_pos), key=self.with_pos_embed(memory, pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0]
#Add and Norm tgt = tgt + self.dropout2(tgt2) tgt = self.norm2(tgt)
#FFN tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt)))) tgt = tgt + self.dropout3(tgt2) tgt = self.norm3(tgt) return tgt
def forward_pre(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): tgt2 = self.norm1(tgt) q = k = self.with_pos_embed(tgt2, query_pos) tgt2 = self.self_attn(q, k, value=tgt2, attn_mask=tgt_mask, key_padding_mask=tgt_key_padding_mask)[0] tgt = tgt + self.dropout1(tgt2) tgt2 = self.norm2(tgt) tgt2 = self.multihead_attn(query=self.with_pos_embed(tgt2, query_pos), key=self.with_pos_embed(memory, pos), value=memory, attn_mask=memory_mask, key_padding_mask=memory_key_padding_mask)[0] tgt = tgt + self.dropout2(tgt2) tgt2 = self.norm3(tgt) tgt2 = self.linear2(self.dropout(self.activation(self.linear1(tgt2)))) tgt = tgt + self.dropout3(tgt2) return tgt
def forward(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): if self.normalize_before: return self.forward_pre(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos) return self.forward_post(tgt, memory, tgt_mask, memory_mask, tgt_key_padding_mask, memory_key_padding_mask, pos, query_pos)

有了一个Decoder Layer的定义,再看Transformer的整个Decoder:

class TransformerDecoder(nn.Module):

#值得注意的是:在使用TransformerDecoder时需要传入的参数有:# tgt:Decoder的输入,memory:Encoder的输出,pos:Encoder的位置编码的输出,query_pos:Object Queries,一堆mask def forward(self, tgt, memory, tgt_mask: Optional[Tensor] = None, memory_mask: Optional[Tensor] = None, tgt_key_padding_mask: Optional[Tensor] = None, memory_key_padding_mask: Optional[Tensor] = None, pos: Optional[Tensor] = None, query_pos: Optional[Tensor] = None): output = tgt
intermediate = []
for layer in self.layers: output = layer(output, memory, tgt_mask=tgt_mask, memory_mask=memory_mask, tgt_key_padding_mask=tgt_key_padding_mask, memory_key_padding_mask=memory_key_padding_mask, pos=pos, query_pos=query_pos) if self.return_intermediate: intermediate.append(self.norm(output))
if self.norm is not None: output = self.norm(output) if self.return_intermediate: intermediate.pop() intermediate.append(output)
if self.return_intermediate: return torch.stack(intermediate)
return output.unsqueeze(0)

然后是把Encoder和Decoder拼在一起,即总的Transformer结构的实现:
此处考虑到字数限制,省略了代码。

实现了Transformer,还剩后面的FFN:

class MLP(nn.Module): """ Very simple multi-layer perceptron (also called FFN)"""
代码略,简单的Pytorch定义layer。

匈牙利匹配HungarianMatcher类:
这个类的目的是计算从targets到predictions的一种最优排列。
predictions比targets的数量多,但我们要进行1-to-1 matching,所以多的predictions将与 匹配。
这个函数整体在构建(13)式,cost_class,cost_bbox,cost_giou,对应的就是(13)式中的几个损失函数,它们的维度都是(b,100,m)。
m包含了这个batch内部所有的 

# pred_logits:[b,100,92]# pred_boxes:[b,100,4]# targets是个长度为b的list,其中的每个元素是个字典,共包含:labels-长度为(m,)的Tensor,元素是标签;boxes-长度为(m,4)的Tensor,元素是Bounding Box。# detr分类输出,num_queries=100,shape是(b,100,92) 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] = [100b, 92] out_bbox = outputs["pred_boxes"].flatten(0, 1) # [batch_size * num_queries, 4] = [100b, 4]
# 准备分类target shape=(m,)里面存储的是类别索引,m包括了整个batch内部的所有gt bbox # Also concat the target labels and boxes tgt_ids = torch.cat([v["labels"] for v in targets])# (m,)[3,6,7,9,5,9,3]# 准备bbox target shape=(m,4),已经归一化了 tgt_bbox = torch.cat([v["boxes"] for v in targets])# (m,4)
#(100b,92)->(100b, m),对于每个预测结果,把目前gt里面有的所有类别值提取出来,其余值不需要参与匹配#对应上述公式,类似于nll loss,但是更加简单 # 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.#行:取每一行;列:只取tgt_ids对应的m列 cost_class = -out_prob[:, tgt_ids]# (100b, m)
# Compute the L1 cost between boxes, 计算out_bbox和tgt_bbox两两之间的l1距离 (100b, m) cost_bbox = torch.cdist(out_bbox, tgt_bbox, p=1)# (100b, m)
# Compute the giou cost betwen boxes, 额外多计算一个giou loss (100b, m) cost_giou = -generalized_box_iou(box_cxcywh_to_xyxy(out_bbox), box_cxcywh_to_xyxy(tgt_bbox))
#得到最终的广义距离(100b, m),距离越小越可能是最优匹配 # Final cost matrix C = self.cost_bbox * cost_bbox + self.cost_class * cost_class + self.cost_giou * cost_giou#(100b, m)--> (b, 100, m) C = C.view(bs, num_queries, -1).cpu()
#计算每个batch内部有多少物体,后续计算时候按照单张图片进行匹配,没必要batch级别匹配,徒增计算 sizes = [len(v["boxes"]) for v in targets]#匈牙利最优匹配,返回匹配索引#enumerate(C.split(sizes, -1))]:(b,100,image1,image2,image3,...) 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]

在得到匹配关系后算loss就水到渠成了。loss_labels计算分类损失,loss_boxes计算回归损失,包含 

PS:作者将继续更新Section2和Section3,请保持关注~

参考文献:

code:
https://github.com/jadore801120/attention-is-all-you-need-pytorch
https://github.com/lucidrains/vit-pytorch
https://github.com/facebookresearch/detr

video:
https://www.bilibili.com/video/av71295187/?spm_id_from=333.788.videocard.8

blog:
https://baijiahao.baidu.com/s?id%3D1651219987457222196&wfr=spider%26for=pc
https://zhuanlan.zhihu.com/p/308301901
https://blog.csdn.net/your_answer/article/details/79160045

◎作者档案

作者:科技猛兽欢迎大家联系极市小编(微信ID:fengcall19)加入极市原创作者行列


推荐阅读


添加极市小助手微信(ID : cvmart2),备注:姓名-学校/公司-研究方向-城市(如:小极-北大-目标检测-深圳),即可申请加入极市目标检测/图像分割/工业检测/人脸/医学影像/3D/SLAM/自动驾驶/超分辨率/姿态估计/ReID/GAN/图像增强/OCR/视频理解等技术交流群:每月大咖直播分享、真实项目需求对接、求职内推、算法竞赛、干货资讯汇总、与 10000+来自港科大、北大、清华、中科院、CMU、腾讯、百度等名校名企视觉开发者互动交流~

△长按添加极市小助手

△长按关注极市平台,获取最新CV干货

觉得有用麻烦给个在看啦~  

您可能也对以下帖子感兴趣

文章有问题?点此查看未经处理的缓存