Transformer源码详解（Pytorch版本）

Transformer源码详解（Pytorch版本）

Pytorch版代码链接如下

GitHub - harvardnlp/annotated-transformer: An annotated implementation of the Transformer paper.

首先来看看attention函数，该函数实现了Transformer中的多头自注意力机制的计算过程。

def attention(query, key, value, mask=None, dropout=None):
​
    d_k = query.size(-1)
    scores = torch.matmul(query, key.transpose(-2, -1))/math.sqrt(d_k)
    if mask is not None:
        scores = scores.masked_fill(mask == 0, -1e9)##使用mask范围内的值进行填充
    p_attn = F.softmax(scores, dim = -1)
    if dropout is not None:
        p_attn = dropout(p_attn)
    return torch.matmul(p_attn, value), p_attn

size(-1)是返回矩阵最后一维的大小。在该函数中，query是四维矩阵（后面会解释），所以返回query的列数，即输入序列的字符的embedding长度。

然后就是query与key的矩阵相乘，再除以一个标量。不过transpose(-2,-1)我觉得有点问题。transpose函数是个矩阵转置函数，比如一个二维矩阵，transpose(1,0)的意思就是将原矩阵第二维放在第一维，第一维放在第二维。再比如一个三维矩阵，假设一开始的矩阵维度分别是x、y、z，transpose(1,0,2)后矩阵维度就变成了y、x、z。大家应该能理解了把。

然后就是softmax，得到的矩阵与value矩阵相乘，最后返回相乘后的矩阵以及p_attn，在后文中我统一称作权重矩阵与P矩阵。

多头自注意力机制类的实现。

class MultiHeadedAttention(nn.Module):
    def __init__(self, h, d_model, dropout=0.1):
        super(MultiHeadedAttention, self).__init__()
        assert d_model % h == 0                             
        self.d_k = d_model // h
        self.h = h
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)
​
    def forward(self, query, key, value, mask=None):
        if mask is not None:
            # Same mask applied to all h heads.
            mask = mask.unsqueeze(1)
        nbatches = query.size(0)
        # 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))]
        # 2) Apply attention on all the projected vectors in batch.
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        # 3) "Concat" using a view and apply a final linear.
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
        return self.linears[-1](x)

首先判断d-model%h是否有余数，如果有就报错。

d-model一般为512（序列符号的embedding长度），h是头数一般为8，两者相除得到d_k的长度，即query、key矩阵中的列数。这个并没有什么原理在里面，只是原论文是这么做的。

def clones(module, N):
    return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])

clones函数的功能是克隆层，不过这些克隆出来的层之间的参数是不共享的。有关ModuleList的讲解可以看这篇文章，非常清晰。

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))]

forward函数中这句代码是关键，也比较难懂，就也就是形成Q、K、V矩阵的过程。

该代码的意思是将传入的Q、K、V三维矩阵经过一层全连接层后，重塑为四维矩阵，并且四维矩阵的第二三维转置。图示可以很好的解释数据结构：

调用attention函数，返回权重矩阵与P矩阵。attention函数具体细节前面已经说过了，在这里我再强调一遍数据的流转细节（我一直认为，看深度学习模型源码，最重要的就是搞懂数据流转过程中维度的变化，可以很好的掌握算法的细节以及体会到并行运算的高效。）

其中F*F就可以理解为每个字符（特征）之间两两相似度分数，即P矩阵。然后P矩阵Softmax后每一行之和等于1，再与Value矩阵相乘，得到权重矩阵。

权重矩阵的每一行可以理解为经过自注意力机制后的每个字符的embedding，即经过考虑其他所有字符影响的embedding。

contiguous的作用是断掉拷贝之间的联系。forward函数最后又把矩阵转置回来，然后经过一层全连接最后返回出去。

class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)
​
    def forward(self, x):
        return self.w_2(self.dropout(F.relu(self.w_1(x))))

用来对解码器的输入端做预处理

class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        self.lut = nn.Embedding(vocab, d_model)
        self.d_model = d_model
​
    def forward(self, x):
        return self.lut(x) * math.sqrt(self.d_model)
class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(0, max_len).unsqueeze(1)
        div_term = torch.exp(torch.arange(0, d_model, 2) *
                             -(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):
        x = x + Variable(self.pe[:, :x.size(1)], requires_grad=False)
        return self.dropout(x)
'''
self.tensor：类成员变量并不会在我们的model.state_dict()，即无法保存；；类成员变量并不会随着model.cuda()复制到gpu上，只能在cpu上。
register_buffer回顾模型保存：torch.save(model.state_dict())，model.state_dict()是一个字典，里边存着我们模型        各个部分的参数。在model中，我们需要更新其中的参数，训练结束将参数保存下来。但在某些时候，我们可能希望模型中的某些参数参        数不更新（从开始到结束均保持不变），但又希望参数保存下来（model.state_dict()），这是我们就会用到register_buffer() 。
不更新，但是不算是模型中的参数（model.state_dict()）
模型保存下来的参数有两种：一种是需要更新的Parameter，另一种是不需要更新的buffer。：
​
'''

PositionalEncoding是编码器位置向量的编码，仔细说一下。弄懂这个我还是强调一点，搞清楚矩阵的维度变化。

首先pe是【5000 512】，position是【5000 1】，div-term是【1 256】。所以position div-term是【5000 256】,经过sin、cos后分别赋值给pe的偶数列与奇数列。pe经过unsqueeze后pe变成了【1 5000 512】。

在Transformer模型中，d_model是模型维度的参数，它指的是模型内部表示（即隐藏层）的大小。这个维度与模型的词嵌入（embeddings）维度相同，并且在整个模型中保持一致，包括编码器（Encoder）和解码器（Decoder）的所有层。

现在我们已经看完了多头自注意力机制、向量Embedding、位置编码Embedding组件。接下来我们看看最后一个组件加和归一化。

class LayerNorm(nn.Module):
    def __init__(self, features, eps=1e-6):
        super(LayerNorm, self).__init__()
        self.a_2 = nn.Parameter(torch.ones(features))
        self.b_2 = nn.Parameter(torch.zeros(features))
        self.eps = eps
​
    def forward(self, x):
        mean = x.mean(-1, keepdim=True)
        std = x.std(-1, keepdim=True)
        return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
class SublayerConnection(nn.Module):
​
    def __init__(self, size, dropout):
        super(SublayerConnection, self).__init__()
        self.norm = LayerNorm(size)
        self.dropout = nn.Dropout(dropout)
​
    def forward(self, x, sublayer):
        return x + self.dropout(sublayer(self.norm(x)))

加和层。

接下来的代码没有什么计算逻辑，都是一些架构上的逻辑。还剩下五个类：EncoderLayer、DecoderLayer、Encoder、Decoder、EncoderDecoder。前两个分别是搭建编码器、解码器中具体的某一层的；第三、第四就是编码器、解码器，其实就是不停复制前两个层，中间穿插标准化层；最后一个就是把编码器、解码器综合在一起的整个Transformers模型类。

class EncoderLayer(nn.Module):
    def __init__(self, size, self_attn, feed_forward, dropout):
        super(EncoderLayer, self).__init__()
        self.self_attn = self_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 2)
        self.size = size
​
    def forward(self, x, mask):
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
        return self.sublayer[1](x, self.feed_forward)

class Encoder(nn.Module):
    def __init__(self, layer, N):
        super(Encoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
​
    def forward(self, x, mask):
        for layer in self.layers:
            x = layer(x, mask)
        return self.norm(x)

这是编码器类，循环调用具体的编码层，层与层之间做归一化。

class DecoderLayer(nn.Module):
    def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
        super(DecoderLayer, self).__init__()
        self.size = size
        self.self_attn = self_attn
        self.src_attn = src_attn
        self.feed_forward = feed_forward
        self.sublayer = clones(SublayerConnection(size, dropout), 3)
​
    def forward(self, x, memory, src_mask, tgt_mask):
        m = memory
        x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
        x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
        return self.sublayer[2](x, self.feed_forward)

class Decoder(nn.Module):
    "Generic N layer decoder with masking."
​
    def __init__(self, layer, N):
        super(Decoder, self).__init__()
        self.layers = clones(layer, N)
        self.norm = LayerNorm(layer.size)
    def forward(self, x, memory, src_mask, tgt_mask):
        for layer in self.layers:
            x = layer(x, memory, src_mask, tgt_mask)
        return self.norm(x)

解码器类与编码器类很相似，就是循环调用解码层，并在层与层之间做归一化。不过解码器的forward函数多了一个memory参数，这个参数是为了接收来自编码器的矩阵。

class EncoderDecoder(nn.Module):
    def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
        super(EncoderDecoder, self).__init__()
        self.encoder = encoder
        self.decoder = decoder
        self.src_embed = src_embed
        self.tgt_embed = tgt_embed
        self.generator = generator
​
    def forward(self, src, tgt, src_mask, tgt_mask):
        return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
​
    def encode(self, src, src_mask):
        return self.encoder(self.src_embed(src), src_mask)
​
    def decode(self, memory, src_mask, tgt, tgt_mask):
        return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)