搞懂 Vision Transformer 原理和代码,看这篇技术综述就够了
极市导读
本文对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 Encoding2 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相似的过程,又输出一排结果,使用黄色的小圆点表示。
所以,用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所做的事情。
所以重点是:我们有一种新的layer,叫self-attention,它的输入和输出和RNN是一模一样的,输入一个sequence,输出一个sequence,它的每一个输出 都看过了整个的输入sequence,这一点与bi-directional RNN相同。但是神奇的地方是:它的每一个输出 可以并行化计算。
1.2 Self-attention:
那么self-attention具体是怎么做的呢?
首先假设我们的input是图4的 ,是一个sequence,每一个input (vector)先乘上一个矩阵 得到embedding,即向量 。接着这个embedding进入self-attention层,每一个向量 分别乘上3个不同的transformation matrix ,以向量 为例,分别得到3个不同的向量 。
接下来使用每个query 去对每个key 做attention,attention就是匹配这2个向量有多接近,比如我现在要对 和 做attention,我就可以把这2个向量做scaled inner product,得到 。接下来你再拿 和 做attention,得到 ,你再拿 和 做attention,得到 ,你再拿 和 做attention,得到 。那这个scaled inner product具体是怎么计算的呢?
式中, 是 跟 的维度。因为 的数值会随着dimension的增大而增大,所以要除以 的值,相当于归一化的效果。
接下来要做的事如图6所示,把计算得到的所有 值取 操作。
取完 操作以后,我们得到了 ,我们用它和所有的 值进行相乘。具体来讲,把 乘上 ,把 乘上 ,把 乘上 ,把 乘上 ,把结果通通加起来得到 ,所以,今天在产生 的过程中用了整个sequence的资讯 (Considering the whole sequence)。如果要考虑local的information,则只需要学习出相应的 , 就不再带有那个对应分支的信息了;如果要考虑global的information,则只需要学习出相应的 , 就带有全部的对应分支的信息了。
同样的方法,也可以计算出 ,如下图8所示, 就是拿query 去对其他的 做attention,得到 ,再与value值 相乘取weighted sum得到的。
经过了以上一连串计算,self-attention layer做的事情跟RNN是一样的,只是它可以并行的得到layer输出的结果,如图9所示。现在我们要用矩阵表示上述的计算过程。
首先输入的embedding是 ,然后用 乘以transformation matrix 得到 ,它的每一列代表着一个vector 。同理,用 乘以transformation matrix 得到 ,它的每一列代表着一个vector 。用 乘以transformation matrix 得到 ,它的每一列代表着一个vector 。
接下来是 与 的attention过程,我们可以把vector 横过来变成行向量,与列向量 做内积,这里省略了 。这样, 就成为了 的矩阵,它由4个行向量拼成的矩阵和4个列向量拼成的矩阵做内积得到,如图11所示。
在得到 以后,如上文所述,要得到 , 就要使用 分别与 相乘再求和得到,所以 要再左乘 矩阵。
到这里你会发现这个过程可以被表示为,如图12所示:输入矩阵 分别乘上3个不同的矩阵 得到3个中间矩阵 。它们的维度是相同的。把 转置之后与 相乘得到Attention矩阵 ,代表每一个位置两两之间的attention。再将它取 操作得到 ,最后将它乘以 矩阵得到输出vector 。
1.3 Multi-head Self-attention:
还有一种multi-head的self-attention,以2个head的情况为例:由 生成的 进一步乘以2个转移矩阵变为 和 ,同理由 生成的 进一步乘以2个转移矩阵变为 和 ,由 生成的 进一步乘以2个转移矩阵变为 和 。接下来 再与 做attention,得到weighted sum的权重 ,再与 做weighted sum得到最终的 。同理得到 。现在我们有了 和 ,可以把它们concat起来,再通过一个transformation matrix调整维度,使之与刚才的 维度一致(这步如图13所示)。
从下图14可以看到 Multi-Head Attention 包含多个 Self-Attention 层,首先将输入 分别传递到 2个不同的 Self-Attention 中,计算得到 2 个输出结果。得到2个输出矩阵之后,Multi-Head Attention 将它们拼接在一起 (Concat),然后传入一个Linear层,得到 Multi-Head Attention 最终的输出 。可以看到 Multi-Head Attention 输出的矩阵 与其输入的矩阵 的维度是一样的。
这里有一组Multi-head Self-attention的解果,其中绿色部分是一组query和key,红色部分是另外一组query和key,可以发现绿色部分其实更关注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所示:
具体的做法是:给每一个位置规定一个表示位置信息的向量 ,让它与 加在一起之后作为新的 参与后面的运算过程,但是这个向量 是由人工设定的,而不是神经网络学习出来的。每一个位置都有一个不同的 。
那到这里一个自然而然的问题是:为什么是 与 相加?为什么不是concatenate?加起来以后,原来表示位置的资讯不就混到 里面去了吗?不就很难被找到了吗?
这里提供一种解答这个问题的思路:
如图15所示,我们先给每一个位置的 append一个one-hot编码的向量 ,得到一个新的输入向量 ,这个向量作为新的输入,乘以一个transformation matrix 。那么:
所以,与 相加就等同于把原来的输入 concat一个表示位置的独热编码 ,再做transformation。
这个与位置编码乘起来的矩阵 是手工设计的,如图17所示。
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取代掉。
2 Transformer的实现和代码解读
2.1 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。
其中,Batch Normalization和Layer Normalization的对比可以概括为图20,Batch Normalization强行让一个batch的数据的某个channel的 ,而Layer Normalization让一个数据的所有channel的 。
接着是一个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操作之前。
因为在翻译的过程中是顺序翻译的,即翻译完第 个单词,才可以翻译第 个单词。通过 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>"。
注意这里transformer模型训练和测试的方法不同:
测试时:
输入<Begin>,解码器输出 I 。 输入前面已经解码的<Begin>和 I,解码器输出have。 输入已经解码的<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的输出结果 。
第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所示。
如下图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所示。
网络的主要组成是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所示,以下分别解释这四个部分:
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到底是怎样变化的,这里我用一个图表达,终结这个问题。
所以,了解了DETR的位置编码之后,你应该明白了其实input embedding和位置编码维度其实是一样的,只是论文图示为了突出二位编码所以画的不一样罢了,如下图所示:
另一点不同的是,原版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不做任何处理。
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主要有两个输入:
Transformer Encoder输出的Embedding与 position encoding 之和。 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没有位置编码:
得到了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的Key和Value就是编码器输出的编码向量(如图30所示),而Query就是Object queries,就是我们的100个格子。
Query可以视作代表不同 的信息,而Key和Value可以视作代表图像的全局信息。
现在通过注意力模块将Query和Key计算,然后加权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. 性能对比:
2. 编码器层数对比实验:
可以发现,编码器层数越多越好,最后就选择6。
下图34为最后一个Encoder Layer的attention可视化,Encoder已经分离了instances,简化了Decoder的对象提取和定位。
3. 解码器层数对比实验:
可以发现,性能随着解码器层数的增加而提升,DETR本不需要NMS,但是作者也进行了,上图中的NMS操作是指DETR的每个解码层都可以输入无序集合,那么将所有解码器无序集合全部保留,然后进行NMS得到最终输出,可以发现性能稍微有提升,特别是AP50。这可以通过以下事实来解释:Transformer的单个Decoder Layer不能计算输出元素之间的任何互相关,因此它易于对同一对象进行多次预测。在第2个和随后的Decoder Layer中,self-attention允许模型抑制重复预测。所以NMS带来的改善随着Decoder Layer的增加而减少。在最后几层,作者观察到AP的一个小损失,因为NMS错误地删除了真实的positive prediction。
类似于可视化编码器注意力,作者在图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