The Annotated Transformer
1 词嵌入
1.1 embeddings
class Embeddings(nn.Module):
def __init__(self,d_model,vocab):
#d_model=512, vocab=当前语言的词表大小
super(Embeddings,self).__init__()
self.lut=nn.Embedding(vocab,d_model)
# one-hot转词嵌入，这里有一个待训练的矩阵E，大小是vocab*d_model
self.d_model=d_model # 512
def forward(self,x):
# x ~ (batch.size, sequence.length, one-hot),
#one-hot大小=vocab，当前语言的词表大小
return self.lut(x)*math.sqrt(self.d_model)
# 得到的10*512词嵌入矩阵，主动乘以sqrt(512)=22.6，
#这里我做了一些对比，感觉这个乘以sqrt(512)没啥用… 求反驳。
#这里的输出的tensor大小类似于(batch.size, sequence.length, 512)
词嵌入矩阵，大小为vocab词个数*d_model词向量长度
1.2 PositionalEncoding
class PositionalEncoding(nn.Module):
"Implement the PE function."
def __init__(self, d_model, dropout, max_len=5000):
#d_model=512,dropout=0.1,
#max_len=5000代表事先准备好长度为5000的序列的位置编码，其实没必要，
#一般100或者200足够了。
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
# Compute the positional encodings once in log space.
pe = torch.zeros(max_len, d_model)
#(5000,512)矩阵，保持每个位置的位置编码，一共5000个位置，
#每个位置用一个512维度向量来表示其位置编码
position = torch.arange(0, max_len).unsqueeze(1)
# (5000) -> (5000,1)
div_term = torch.exp(torch.arange(0, d_model, 2) *
-(math.log(10000.0) / d_model))
# (0,2,…, 4998)一共准备2500个值，供sin, cos调用
pe[:, 0::2] = torch.sin(position * div_term) # 偶数下标的位置
pe[:, 1::2] = torch.cos(position * div_term) # 奇数下标的位置
pe = pe.unsqueeze(0)
# (5000, 512) -> (1, 5000, 512) 为batch.size留出位置
self.register_buffer('pe', pe)
def forward(self, x):
x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
# 接受1.Embeddings的词嵌入结果x，
#然后把自己的位置编码pe，封装成torch的Variable(不需要梯度)，加上去。
#例如，假设x是(30,10,512)的一个tensor，
#30是batch.size, 10是该batch的序列长度, 512是每个词的词嵌入向量；
#则该行代码的第二项是(1, min(10, 5000), 512)=(1,10,512)，
#在具体相加的时候，会扩展(1,10,512)为(30,10,512)，
#保证一个batch中的30个序列，都使用（叠加）一样的位置编码。
return self.dropout(x) # 增加一次dropout操作
# 注意，位置编码不会更新，是写死的，所以这个class里面没有可训练的参数。
该方法为写死的位置编码信息，比起可训练的位置参数，好处是减少参数量提高效率，同时可以计算更长的位置信息（如：长度为5000的句子）。
函数中，pos代表vocab在句子中的位置，i代表512维度中每一维度的数值，该函数能表示位置信息的原因是，每个dimension对应的pos弦波的偏差offset和频率frequency不同。如下图所示：
使用的时候，使用sequential函数，先计算token ebedding，然后计算pe加上去。
3 多头注意力
3.1 attenion
def attention(query, key, value, mask=None, dropout=None):
# query, key, value的形状类似于(30, 8, 10, 64), (30, 8, 11, 64),
#(30, 8, 11, 64)，例如30是batch.size，即当前batch中有多少一个序列；
# 8=head.num，注意力头的个数；
# 10=目标序列中词的个数，64是每个词对应的向量表示；
# 11=源语言序列传过来的memory中，当前序列的词的个数，
# 64是每个词对应的向量表示。
# 类似于，这里假定query来自target language sequence；
# key和value都来自source language sequence.
"Compute 'Scaled Dot Product Attention'"
d_k = query.size(-1) # 64=d_k
scores = torch.matmul(query, key.transpose(-2, -1)) /
math.sqrt(d_k) # 先是(30,8,10,64)和(30, 8, 64, 11)相乘，
#（注意是最后两个维度相乘）得到(30,8,10,11)，
#代表10个目标语言序列中每个词和11个源语言序列的分别的“亲密度”。
#然后除以sqrt(d_k)=8，防止过大的亲密度。
#这里的scores的shape是(30, 8, 10, 11)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
#使用mask，对已经计算好的scores，按照mask矩阵，填-1e9，
#然后在下一步计算softmax的时候，被设置成-1e9的数对应的值~0,被忽视
p_attn = F.softmax(scores, dim = -1)
#对scores的最后一个维度执行softmax，得到的还是一个tensor,
#(30, 8, 10, 11)
if dropout is not None:
p_attn = dropout(p_attn) #执行一次dropout
return torch.matmul(p_attn, value), p_attn
#返回的第一项，是(30,8,10, 11)乘以（最后两个维度相乘）
#value=(30,8,11,64)，得到的tensor是(30,8,10,64)，
#和query的最初的形状一样。另外，返回p_attn，形状为(30,8,10,11).
#注意，这里返回p_attn主要是用来可视化显示多头注意力机制。
1 query和key,value的shape可以不同，在于第三个维度序列中的词数。对于self-attention,qkv是相同的，对于encoder-decoder的attention，q是decoder的hidden_state，kv是encoder的最后输出。这里与gru/lstm+attention的encoder-decoder方案不同。后者是encoder每层的hidden_state与decoder的hidden_state广播后拼接,softmax(linear(tanh(linear(enc_and_dec_states))))*enc_states作为最终结果（加法attention）。前者只取最后的hidden_state，作为decoder self-attention的kv。
2 attention有两种，addtitive attention和dot-product attention。前者是用一个前馈神经网络，用单个隐藏层完成的。点乘更快同时空间效率更高，因为使用了矩阵乘法。
3 文中用的是dot-product attention，同时加上了根号dk作为scaled，原因是当dk大的时候，会导致点乘的结果变得巨大，从而导致进入softmax函数后，会获得特别小的梯度。为了防止因为dk维度的变大导致这个问题，scaled解决dk维度大小对结果的影响。
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
# h=8, d_model=512
"Take in model size and number of heads."
super(MultiHeadedAttention, self).__init__()
assert d_model % h == 0 # We assume d_v always equals d_k 512%8=0
self.d_k = d_model // h # d_k=512//8=64
self.h = h #8
self.linears = clones(nn.Linear(d_model, d_model), 4)
#定义四个Linear networks, 每个的大小是(512, 512)的，
#每个Linear network里面有两类可训练参数，Weights，
#其大小为512*512，以及biases，其大小为512=d_model。
self.attn = None
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):
# 注意，输入query的形状类似于(30, 10, 512)，
# key.size() ~ (30, 11, 512),
#以及value.size() ~ (30, 11, 512)
if mask is not None: # Same mask applied to all h heads.
mask = mask.unsqueeze(1) # mask下回细细分解。
nbatches = query.size(0) #e.g., nbatches=30
# 1) Do all the linear projections in batch from
#d_model => h x d_k
query, key, value = [l(x).view(nbatches, -1, self.h, self.d_k)
.transpose(1, 2) for l, x in
zip(self.linears, (query, key, value))]
# 这里是前三个Linear Networks的具体应用，
#例如query=(30,10, 512) -> Linear network -> (30, 10, 512)
#-> view -> (30,10, 8, 64) -> transpose(1,2) -> (30, 8, 10, 64)
#，其他的key和value也是类似地，
#从(30, 11, 512) -> (30, 8, 11, 64)。
# 2) Apply attention on all the projected vectors in batch.
x, self.attn = attention(query, key, value, mask=mask,
dropout=self.dropout)
#调用上面定义好的attention函数，输出的x形状为(30, 8, 10, 64)；
#attn的形状为(30, 8, 10=target.seq.len, 11=src.seq.len)
# 3) "Concat" using a view and apply a final linear.
x = x.transpose(1, 2).contiguous().
view(nbatches, -1, self.h * self.d_k)
# x ~ (30, 8, 10, 64) -> transpose(1,2) ->
#(30, 10, 8, 64) -> contiguous() and view ->
#(30, 10, 8*64) = (30, 10, 512)
return self.linears[-1](x)
#执行第四个Linear network，把(30, 10, 512)经过一次linear network，
#得到(30, 10, 512).
这里一共有四个矩阵。h=8的情况下，d_k=d_v = d_model//h= 512/8=64。因为并行计算，多头的计算效率和单头的是一样的。attention计算流程：
1 .query.shape(30,10,512) key.shape(30,11,512)value(30,11,512)。qkv分别乘512*512的矩阵，输出的（30,10,512）先通过view转换成8头的形式（30,10,8,64），再通过transpose（1,2）转换成（30,8,10,64），形成最终输入attention的qkv。
2 通过attention函数，输出的形状为(30, 8, 10, 64)，view转成（30,10,512）在经过一个512*512的线形成，得到（30,10,512）。
4.SubLayerConnection 子层连接
所谓SubLayer，指的是两个部分：3.MultiHeadAttention以及4.PositionwiseFeedForward。特别要提的一点是，这两个子层的输入和输出，都是类似于(batch.size, sequence.length, d_model)这样的tensor。
class LayerNorm(nn.Module):
"Construct a layernorm module (See citation for details)."
def __init__(self, features, eps=1e-6):
# features=d_model=512, eps=epsilon 用于分母的非0化平滑
super(LayerNorm, self).__init__()
self.a_2 = nn.Parameter(torch.ones(features))
# a_2 是一个可训练参数向量，(512)
self.b_2 = nn.Parameter(torch.zeros(features))
# b_2 也是一个可训练参数向量, (512)
self.eps = eps
def forward(self, x):
# x 的形状为(batch.size, sequence.len, 512)
mean = x.mean(-1, keepdim=True)
# 对x的最后一个维度，取平均值，得到tensor (batch.size, seq.len)
std = x.std(-1, keepdim=True)
# 对x的最后一个维度，取标准方差，得(batch.size, seq.len)
return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
# 本质上类似于（x-mean)/std，不过这里加入了两个可训练向量
# a_2 and b_2，以及分母上增加一个极小值epsilon，用来防止std为0
# 的时候的除法溢出
这里衍生出来一个问题：为什么transformer要用layer norm而不是batch norm?
一般来说，如果你的特征依赖于不同样本间的统计参数，那BN更有效。因为它抹杀了不同特征之间的大小关系，但是保留了不同样本间的大小关系。（CV领域）
而在NLP领域，LN就更加合适。因为它抹杀了不同样本间的大小关系，但是保留了一个样本内不同特征之间的大小关系。对于NLP或者序列任务来说，一条样本的不同特征，其实就是时序上字符取值的变化，样本内的特征关系是非常紧密的。
class SublayerConnection(nn.Module):
"""
A residual connection followed by a layer norm.
Note for code simplicity the norm is first as opposed to last.
"""
def __init__(self, size, dropout):
super(SublayerConnection, self).__init__()
self.norm = LayerNorm(size)
self.dropout = nn.Dropout(dropout)
def forward(self, x, sublayer):
"Apply residual connection to any sublayer with the same size."
return x + self.dropout(sublayer(self.norm(x)))
SubLayerConnection类主要实现两个功能，残差Add以及Norm（使用上面的LayerNorm类。
x + self.dropout(sublayer(self.norm(x)))实现残差功能。
残差首先解决梯度消失为题。梯度消失是犹豫反向传播时，一旦其中某一个导数很小，多次连乘后梯度可能越来越小，这就是常说的梯度消散，对于深层网络，传到浅层几乎就没了。但是如果使用了残差，每一个导数就加上了一个恒等项1，dh/dx=d(f+x)/dx=1+df/dx。此时就算原来的导数df/dx很小，这时候误差仍然能够有效的反向传播，这就是核心思想。
其次，经网络的退化才是难以训练深层网络根本原因所在，而不是梯度消散。虽然梯度范数大，但是如果网络的可用自由度对这些范数的贡献非常不均衡，也就是每个层中只有少量的隐藏单元对不同的输入改变它们的激活值，而大部分隐藏单元对不同的输入都是相同的反应，此时整个权重矩阵的秩不高。并且随着网络层数的增加，连乘后使得整个秩变的更低。这也是我们常说的网络退化问题，虽然是一个很高维的矩阵，但是大部分维度却没有信息，表达能力没有看起来那么强大。残差连接正是强制打破了网络的对称性。
5.PositionwiseFeedForward
class PositionwiseFeedForward(nn.Module):
"Implements FFN equation."
def __init__(self, d_model, d_ff, dropout=0.1):
# d_model = 512
# d_ff = 2048 = 512*4
super(PositionwiseFeedForward, self).__init__()
self.w_1 = nn.Linear(d_model, d_ff)
# 构建第一个全连接层，(512, 2048)，其中有两种可训练参数：
# weights矩阵，(512, 2048)，以及
# biases偏移向量, (2048)
self.w_2 = nn.Linear(d_ff, d_model)
# 构建第二个全连接层, (2048, 512)，两种可训练参数：
# weights矩阵，(2048, 512)，以及
# biases偏移向量, (512)
self.dropout = nn.Dropout(dropout)
def forward(self, x):
# x shape = (batch.size, sequence.len, 512)
# 例如, (30, 10, 512)
return self.w_2(self.dropout(F.relu(self.w_1(x))))
# x (30, 10, 512) -> self.w_1 -> (30, 10, 2048)
# -> relu -> (30, 10, 2048)
# -> dropout -> (30, 10, 2048)
# -> self.w_2 -> (30, 10, 512)是输出的shape
FFN(x)=max(0,xW1+b1)W2+b2。实现self-attention后的全连接层。
6.EncoderLayer
class EncoderLayer(nn.Module):
"Encoder is made up of self-attn and "
"feed forward (defined below)"
def __init__(self, size, self_attn, feed_forward, dropout):
# size=d_model=512
# self_attn = MultiHeadAttention对象, first sublayer
# feed_forward = PositionwiseFeedForward对象，second sublayer
# dropout = 0.1 (e.g.)
super(EncoderLayer, self).__init__()
self.self_attn = self_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 2)
# 使用深度克隆方法，完整地复制出来两个SublayerConnection
self.size = size # 512
def forward(self, x, mask):
"Follow Figure 1 (left) for connections."
# x shape = (30, 10, 512)
# mask 是(batch.size, 10,10)的矩阵，类似于当前一个词w，有哪些词是w可见的
# 源语言的序列的话，所有其他词都可见，除了"<blank>"这样的填充；
# 目标语言的序列的话，所有w的左边的词，都可见。
x = self.sublayer[0](x,
lambda x: self.self_attn(x, x, x, mask))
# x (30, 10, 512) -> self_attn (MultiHeadAttention)
# shape is same (30, 10, 512) -> SublayerConnection
# -> (30, 10, 512)
return self.sublayer[1](x, self.feed_forward)
# x 和feed_forward对象一起，给第二个SublayerConnection
这里给self_attention传了一个mask矩阵。矩阵大小为（batch_size,seq_len,seq_len）。对于encoder而言，仅有padding的部分是不可见的，需要mask掉。对于decoder而言，需要mask掉矩阵的上半部分，因为encoder的时候positon i的词只能依赖于i以及i以前的outputs。
对于参数量：
其一，一个MultiHeadAttention对象，里面有4个Linear network；参数个数是：4(512512+512)=1,050,624
其二，一个PositionwiseFeedForward对象，里面有2个Linear network；参数个数是：2,099,712。
其三，两个SublayerConnection对象，每个对象使用一个LayerNorm对象，LayerNorm里面有两个可训练参数向量a_2和b_2。则，可训练参数个数是：
2 * (512 + 512) = 2048。
综合上面三个部分：一共有: 1,050,624 + 2,099,712 + 2048 = 3,152,384个参数。（一个EncoderLayer)
7. Encoder
class Encoder(nn.Module):
"Core encoder is a stack of N layers"
def __init__(self, layer, N):
# layer = one EncoderLayer object, N=6
super(Encoder, self).__init__()
self.layers = clones(layer, N)
# 深copy，N=6，
self.norm = LayerNorm(layer.size)
# 定义一个LayerNorm，layer.size=d_model=512
# 其中有两个可训练参数a_2和b_2
def forward(self, x, mask):
"Pass the input (and mask) through each layer in turn."
# x is alike (30, 10, 512)
# (batch.size, sequence.len, d_model)
# mask是类似于(batch.size, 10, 10)的矩阵
for layer in self.layers:
x = layer(x, mask)
# 进行六次EncoderLayer操作
return self.norm(x)
# 最后做一次LayerNorm，最后的输出也是(30, 10, 512) shape
数数这里的可训练参数的个数：
6层EncoderLayers: 每一层是 3,152,384个参数；
LayerNorm里面有两个512维度的向量，有1024个参数；
则一共有：3,152,384*6 + 1,024 = 18,915,328个参数(Encoder类)。
再假设，源语言词表大小是30,000词，则embeddings部分有30,000 * 512 = 15,360,000个参数（词嵌入矩阵）。
这样的话，源语言的Embeddings+Encoder，一共有34,275,328个可训练参数。
8 DecoderLayer
class DecoderLayer(nn.Module):
"Decoder is made of self-attn, src-attn, "
"and feed forward (defined below)"
def __init__(self, size, self_attn, src_attn,
feed_forward, dropout):
# size = d_model=512,
# self_attn = one MultiHeadAttention object，目标语言序列的
# src_attn = second MultiHeadAttention object, 目标语言序列
# 和源语言序列之间的
# feed_forward 一个全连接层
# dropout = 0.1
super(DecoderLayer, self).__init__()
self.size = size # 512
self.self_attn = self_attn
self.src_attn = src_attn
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 3)
# 需要三个SublayerConnection, 分别在
# self.self_attn, self.src_attn, 和self.feed_forward
# 的后边
def forward(self, x, memory, src_mask, tgt_mask):
"Follow Figure 1 (right) for connections."
m = memory # (batch.size, sequence.len, 512)
# 来自源语言序列的Encoder之后的输出，作为memory
# 供目标语言的序列检索匹配：（类似于alignment in SMT)
x = self.sublayer[0](x,
lambda x: self.self_attn(x, x, x, tgt_mask))
# 通过一个匿名函数，来实现目标序列的自注意力编码
# 结果扔给sublayer[0]:SublayerConnection
x = self.sublayer[1](x,
lambda x: self.src_attn(x, m, m, src_mask))
# 通过第二个匿名函数，来实现目标序列和源序列的注意力计算
# 结果扔给sublayer[1]:SublayerConnection
return self.sublayer[2](x, self.feed_forward)
# 走一个全连接层，然后
# 结果扔给sublayer[2]:SublayerConnection
这里估计一下DecoderLayer的可训练参数的个数：
其一，两个MultiHeadAttention对象，每个对象有1,050,624个参数；
其二，feedforward对象，2,099,712个参数；
其三，三个SublayerConnection, 3*(512+512) = 3,072个参数。
综合起来有，2*1,050,624 + 2,099,712 + 3,072 = 4,204,032个参数 (DecoderLayer)。
class Decoder(nn.Module):
"Generic N layer decoder with masking."
def __init__(self, layer, N):
# layer = DecoderLayer object
# N = 6
super(Decoder, self).__init__()
self.layers = clones(layer, N)
# 深度copy六次DecoderLayer
self.norm = LayerNorm(layer.size)
# 初始化一个LayerNorm
def forward(self, x, memory, src_mask, tgt_mask):
for layer in self.layers:
x = layer(x, memory, src_mask, tgt_mask)
# 执行六次DecoderLayer
return self.norm(x)
# 执行一次LayerNorm
Decoder类的可训练参数的个数是：
其一，六层DecoderLayer: 6*4,204,032；
其二，一个LayerNorm, 1,024个参数；
则一共有：6*4,204,032 + 1024 = 25,225,216个参数。
如果假设目标语言也是30,000的词表大小，则又增加30,000 * 512=15,360,000个参数，从而有：40,585,216个可训练参数。
9 generator
class Generator(nn.Module):
"Define standard linear + softmax generation step."
def __init__(self, d_model, vocab):
# d_model=512
# vocab = 目标语言词表大小
super(Generator, self).__init__()
self.proj = nn.Linear(d_model, vocab)
# 定义一个全连接层，可训练参数个数是(512 * trg_vocab_size) +
# trg_vocab_size
def forward(self, x):
return F.log_softmax(self.proj(x), dim=-1)
# x 类似于 (batch.size, sequence.length, 512)
# -> proj 全连接层 (30, 10, trg_vocab_size) = logits
# 对最后一个维度执行log_soft_max
# 得到(30, 10, trg_vocab_size)
上面类里面有一个全连接层，参数个数是，一个权重矩阵 512*trg_vocab_size，一个偏移向量 (trg_vocab_size)。
如果假设trg_vocab_size=30,000, 则可训练参数的个数是：
512*30000 + 30000 = 15,390,000个 (Generator)。
到目前位置，我们基本阐述完毕Encoder, Decoder，Generator，以及source sequence and target sequence的Embedding。
10 EncoderDeocder
class EncoderDecoder(nn.Module):
"""
A standard Encoder-Decoder architecture.
Base for this and many other models.
"""
def __init__(self, encoder, decoder,
src_embed, tgt_embed, generator):
super(EncoderDecoder, self).__init__()
self.encoder = encoder
# Encoder对象
self.decoder = decoder
# Decoder对象
self.src_embed = src_embed
# 源语言序列的编码，包括词嵌入和位置编码
self.tgt_embed = tgt_embed
# 目标语言序列的编码，包括词嵌入和位置编码
self.generator = generator
# 生成器
def forward(self, src, tgt, src_mask, tgt_mask):
"Take in and process masked src and target sequences."
return self.decode(self.encode(src, src_mask), src_mask,
tgt, tgt_mask)
# 先对源语言序列进行编码，
# 结果作为memory传递给目标语言的编码器
def encode(self, src, src_mask):
# src = (batch.size, seq.length)
# src_mask 负责对src加掩码
return self.encoder(self.src_embed(src), src_mask)
# 对源语言序列进行编码，得到的结果为
# (batch.size, seq.length, 512)的tensor
def decode(self, memory, src_mask, tgt, tgt_mask):
return self.decoder(self.tgt_embed(tgt),
memory, src_mask, tgt_mask)
# 对目标语言序列进行编码，得到的结果为
# (batch.size, seq.length, 512)的tensor
上面的类里面，有五大属性，分别的可训练参数的个数是：
（假设源语言和目标语言的词表都是30,000)
源语言的Embeddings+Encoder，一共有34,275,328个可训练参数 (Embedding+Encoder);
假设目标语言也是30,000的词表大小，则又增加30,000 * 512=15,360,000个参数，从而有：40,585,216个可训练参数(Embedding+Decoder);
如果假设trg_vocab_size=30,000, 则可训练参数的个数是：
512*30000 + 30000 = 15,390,000个(Generator).
合计为：90,250,544个参数，9千万。稍稍扩大一下词表大小，轻松一亿个参数。所以，170亿参数，也没啥。
以上就是transformer代码解析部分。后续更新bert及其变种论文解析。
