上一篇DETR代码学习笔记（一） 记录完了DETR的backbone以及数据进入encoder之前的部分，这篇开始介绍DETR的transformer部分，大体上DETR的transformer和Attention is All You Need 中提出框架基本是一致的，只是数据上有些不同。
首先是encoder部分，整个encoder部分，从代码上可以很直观的看出他的整体流程，其中的最初的q，k是由backbone中得到的Feature Map加上位置编码得到，这个位置编码是由backbone中与Feature Map一同输出的mask生成的，如果忘了这个是怎么来的可以回去看上一篇加深印象。同样接上一篇，还是假设输入为[2,768,768]，得到的特征图为[2,256,24,24]，reshape并转换维度之后得到[576,2,256]，此时的初始的q，k相等（Feature Map加上位置编码），而v则是Feature Map，他并没有没有加上位置编码，q，k，v的维度保持不变还是[576,2,256]
进入自注意力层后q，k，v通过线性层进行初始化，之后reshape，由输入的[HWxNxC]->[NXnum_heads,HxW,head_dim]，即[576,2,256]->[16,576,32]，通过如下的公式保持最后的输出维度不变，还是[576,2,256]。
encoder的代码：
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):
# q,k由最初输入的src加上pos的位置编码构成,且q=k,shape为[576,2,256]
# 该位置编码是由backbone中输出的mask生成
q = k = self.with_pos_embed(src, pos)
# 自注意力层,src2 = softmax(q*kt/sqrt(dk))*v,其中dk=32
# 进入自注意力层后会对q,k,v进行reshape,由输入的[HWxNxC]->[NXnum_heads,HxW,head_dim],即[576,2,256]->[16,576,32]
# 自注意力层的输出依旧是[576,2,256]
src2 = self.self_attn(q, k, value=src, attn_mask=src_mask,
key_padding_mask=src_key_padding_mask)[0]
src = src + self.dropout1(src2)
src = self.norm1(src)
src2 = self.linear2(self.dropout(self.activation(self.linear1(src))))
src = src + self.dropout2(src2)
src = self.norm2(src)
return src从代码上能比较直观的得到encoder的框架图。
原文中encoder有6层，也就是第一层encoder的输出作为下一层encoder的输入，直到第六层最后输出的memory，这个memory将作为decoder的输入。
接下来就是decoder部分，decoder与encoder的输入上存在差异，最初自注意力层的输入是词嵌入向量，其中q和k是一个[100,2,256]全零张量加上词嵌入向量，v则是[100,2,256]全零张量，之后进入自注意力层，这里的流程和上面encoder中是一样的。
自注意力层的输出将作为多头注意力层中的q，而k和v来自encoder的输出，其中k还要加上位置编码。此时k和v维度为[576,2,256]，q的维度为[100,2,256]，同样的在计算权重时，会将各个张量进行reshape，即q:[100,2,256]->[16,100,32]，k,v:[576,2,256]->[16,576,32]，再利用公式计算，最后多头注意力层的输出依旧是[100,2,256]。
decoder代码：
class TransformerDecoderLayer(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)
self.multihead_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.norm3 = nn.LayerNorm(d_model)
self.dropout1 = nn.Dropout(dropout)
self.dropout2 = nn.Dropout(dropout)
self.dropout3 = 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, 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):
# q,k由最初输入的tgt加上query_pos的词嵌入向量构成,且q=k,shape为[100,2,256]
# 其中tgt是shape为[100,2,256]的全零输入
q = k = self.with_pos_embed(tgt, query_pos)
# 自注意力层,tgt2 = softmax(q*kt/sqrt(dk))*v,其中dk=32
# 进入自注意力层后会对q,k,v进行reshape,由输入的[HWxNxC]->[NXnum_heads,HxW,head_dim],即[100,2,256]->[16,100,32]
# 自注意力层的输出依旧是[100,2,256]
tgt2 = self.self_attn(q, k, value=tgt, attn_mask=tgt_mask,
key_padding_mask=tgt_key_padding_mask)[0]
tgt = tgt + self.dropout1(tgt2)
tgt = self.norm1(tgt)
# 多头注意力层,计算方式相同tgt2 = softmax(q*kt/sqrt(dk))*v,其中dk=32
# 但是出入的shape发生变化,memory是encoder的输出,shape为[576,2,256],用它作为k,v,k还要加上位置编码
# 多头自注意力层同样对q,k,v进行reshape,由输入的[HWxNxC]->[NXnum_heads,HxW,head_dim]
# 即q:[100,2,256]->[16,100,32]与k,v:[576,2,256]->[16,576,32]
# 多头注意力层的输出依旧是[100,2,256]
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]
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根据代码也比较好得到如下的框架图
与encoder相同，decoder也有6层，第一层decoder的输出作为下一层decoder的输入，最后一层的输出会边界框和类别分类。
将以上的流程组合后得到整体的框架图。
最后附上论文原文中的框架图
到这里大致把transformer部分的处理过程理清，接下来是构建损失函数。有空再整理一波。
