前言
在阅读完《Attention is All You Need》后，我不禁对Transformer架构产生了浓厚的兴趣。于是，跟着原论文，我把Transformer的架构自己敲了一遍。
在阅读完ViT的论文后，我也试图操作一下Transformer在CV领域的工作，项目还是选当时我做CNN的那个花朵分了项目，毕竟考虑到CPU功能是有限的……
这篇微博，是对我初次接触Transformer的一些小感悟。
Transformer架构复现
自注意力
class SelfAttention(nn.Module):
def __init__(self,dropout=0.1):
super(SelfAttention,self).__init__()
self.dropout = nn.Dropout(dropout)
self.softmax = nn.Softmax(dim=-1)这是pytorch中对类的经典初始化，在此我也定义了归一化函数和随机失活率，都是为了过拟合嘛。
def forward(self,Q,K,V,mask=None):
d_k = Q.size(-1)
score = torch.matmul(Q, K.transpose(-2,-1)) / math.sqrt(d_k)
if mask is not None:
score = score.masked_fill(mask==0, float('-inf'))
attn = self.softmax(score)
attn = self.dropout(attn)
out = torch.matmul(attn, V)
return out, attn这是自注意力的前向传播过程，首先就是实现公式
计算每个Q和K的相似程度，然后缩放，softmax处理，最后乘V矩阵。
这里为什么还要有掩码呢？有时候为了格式上的匹配，我们可能在数据中引入没有实际内容的占位元素，这样我们就需要用掩码把它遮住了。
多头自注意力
class MultiHeadAttention(nn.Module):
def __init__(self,d_model,n_heads,dropout=0.1):
super(MultiHeadAttention,self).__init__()
assert d_model % n_heads == 0
self.d_k = d_model // n_heads
self.n_heads = n_heads
self.W_q = nn.Linear(d_model,d_model)
self.W_k = nn.Linear(d_model,d_model)
self.W_v = nn.Linear(d_model,d_model)
self.fc = nn.Linear(d_model,d_model)
self.attention = SelfAttention(dropout)
self.dropout = nn.Dropout(dropout)
self.norm = nn.LayerNorm(d_model)这是多头自注意力的初始化过程，其实就是q，k，v的线性映射，随机失活的作用，全连接层，还有LayerNorm归一化。
def forward(self,q,k,v,mask=None):
batch_size = q.size(0)
Q = self.W_q(q).view(batch_size,-1,self.n_heads,self.d_k).transpose(1,2)
K = self.W_k(k).view(batch_size,-1,self.n_heads,self.d_k).transpose(1,2)
V = self.W_v(v).view(batch_size,-1,self.n_heads,self.d_k).transpose(1,2)
out, attn = self.attention(Q,K,V,mask)
out = out.transpose(1,2).contiguous().view(batch_size,-1,self.n_heads*self.d_k)
out = self.fc(out)
return self.norm(out + q), attn这是多头自注意力的前向传播过程。Q，K，V的运算已经阐述，只需要说一下out。
每一个out经过上一步自注意力的运算，然后用view改变一下形状（运算的时候为了方便，把矩阵第二维度Seq_len和矩阵第三维度Heads_num交换了位置，现在换回来，再把最后两维乘在一起，将四维张量变成三位张量），最后将不同头的out进行全连接，结果依旧要经过残差连接和层归一化。
前馈神经网络
class FeedForward(nn.Module):
def __init__(self,d_model,d_ff,dropout=0.1):
super(FeedForward,self).__init__()
self.fc1 = nn.Linear(d_model,d_ff)
self.fc2 = nn.Linear(d_ff,d_model)
self.dropout = nn.Dropout(dropout)
self.norm = nn.LayerNorm(d_model)这是前馈神经网络的初始化。可以看到，前馈神经网络就是两层线性映射和残差连接，归一化构成。
def forward(self,x):
out = self.fc2(self.dropout(torch.relu(self.fc1(x))))
return self.norm(x + out)其前向传播部分也很简单，不再赘述。就是在两次线性映射时引入激活函数，让其有非线性关系的处理方式。
编码器层
class EncoderLayer(nn.Module):
def __init__(self,d_model,n_heads,d_ff,dropout=0.1):
super(EncoderLayer,self).__init__()
self.self_attn = MultiHeadAttention(d_model,n_heads,dropout)
self.ffn = FeedForward(d_model,d_ff,dropout)编码器，顾名思义，就是把先前写出来的一层多头自注意力编码器用循环的方式输出，而这里先写每个编码器层，这种封装的编程思想确实提供了大大的方便。根据论文的图解，编码器就包括两部分，先经过多头自注意力，然后就是前馈神经网络。
def forward(self,src,src_mask=None):
out,_ = self.self_attn(src,src,src,src_mask)
out = self.ffn(out)
return out没什么好说的，其实就是实现过程。
解码器层
class DecoderLayer(nn.Module):
def __init__(self,d_model,n_heads,d_ff,dropout=0.1):
super(DecoderLayer,self).__init__()
self.self_attn = MultiHeadAttention(d_model,n_heads,dropout)
self.cross_attn = MultiHeadAttention(d_model,n_heads,dropout)
self.ffn = FeedForward(d_model,d_ff,dropout)解码器层，同样的道理，也是为了后期写解码器做铺垫。根据原论文的示意图，解码器层包括掩码自注意力，交叉自注意力，前馈神经网络。这自然就是将其一一实现的过程。
def forward(self,tgt,memory,tgt_mask=None,memory_mask=None):
out,_ = self.self_attn(tgt,tgt,tgt,tgt_mask)
out,_ = self.cross_attn(out,memory,memory,memory_mask)
out = self.ffn(out)
return out没什么好说的，三行代码对应三个过程。
位置编码
class PositionalEncoding(nn.Module):
def __init__(self,d_model,max_len=5000):
super().__init__()
pe = torch.zeros(max_len,d_model)
position = torch.arange(0,max_len,dtype=torch.float32).unsqueeze(1)
div_term = torch.exp(torch.arange(0,d_model,2).float()*(-math.log(10000.0)/d_model))
pe[:,0::2] = torch.sin(position*div_term)
pe[:,1::2] = torch.cos(position*div_term)
pe = pe.unsqueeze(0)
self.register_buffer("pe",pe)
def forward(self,x):
seq_len = x.size(1)
return x + self.pe[:,:seq_len,:]根据论文给的公式，采用三角函数位置编码，通过正弦、余弦函数交替映射奇偶维度：
$$ PE_{pos,2i} = \sin\left(pos / 10000^{2i/d_{model}}\right) $$ $$ PE_{pos,2i+1} = \cos\left(pos / 10000^{2i/d_{model}}\right) $$
将计算得到的位置编码向量，与词嵌入向量逐维相加，融合位置信息与语义信息，一同送入后续注意力层计算。
简单说一下编程中的细节：position扩展维度是为了更好的矩阵运算，而最后我们只需要截取seq_len长度即可了。
self.register_buffer(“pe”, pe) 是 PyTorch 里专门用来注册非训练参数、但需要保存 / 加载 / 和模型一起移动设备的张量的核心方法，这也对应了位置编码不参与训练的原理。
编码器和解码器
class Encoder(nn.Module):
def __init__(self,vocab_size,d_model,d_ff,n_heads,num_layers,max_len=5000,dropout=0.1):
super().__init__()
self.Embedding = nn.Embedding(vocab_size,d_model)
self.pos_encoding = PositionalEncoding(d_model,max_len)
self.layers = nn.ModuleList([EncoderLayer(d_model,n_heads,d_ff,dropout)for _ in range(num_layers)])
def forward(self,src,src_mask=None):
out = self.Embedding(src) * math.sqrt(self.Embedding.embedding_dim)
out = self.pos_encoding(out)
for layer in self.layers:
out = layer(out,src_mask)
return out编码器，实际上就是将编码器层重复几遍，再补充上词语编码，位置编码处理，残差连接，归一化等。开始乘一个数，是为了让特征对比更加明显，这也印证了softmax前为什么要再除以一个数了，防止过大出现梯度消失/梯度爆炸。
class Decoder(nn.Module):
def __init__(self,vocab_size,d_model,d_ff,n_heads,num_layers,max_len=5000,dropout=0.1):
super().__init__()
self.Embedding = nn.Embedding(vocab_size,d_model)
self.pos_encoding = PositionalEncoding(d_model,max_len)
self.layers = nn.ModuleList([DecoderLayer(d_model,n_heads,d_ff,dropout)for _ in range(num_layers)])
self.fc_out = nn.Linear(d_model,vocab_size)
def forward(self,tgt,memory,tgt_mask=None,memory_mask=None):
out = self.Embedding(tgt) * math.sqrt(self.Embedding.embedding_dim)
out = self.pos_encoding(out)
for layer in self.layers:
out = layer(tgt=out, memory=memory, tgt_mask=tgt_mask, memory_mask=memory_mask)
out = self.fc_out(out)
return out解码器，实际上就是将解码器层重复几遍，只不过传入了掩码，交叉自注意力提取了编码器中的Q，其余操作和编码器一致。