图解Transformer

前言

  • Google 提出了 Transformer 模型，用 Self Attention 的结构，取代了以往 NLP 任务中的 RNN 网络结构
  • 优点
    • 使得模型训练过程能够并行计算
  • Transformer 使用了 Seq2Seq任务中常用的结构——包括两个部分：Encoder 和 Decoder。


一、从整体宏观来理解 Transformer

  • Transformer 可以拆分为 2 部分：
    • 编码部分(encoding component)
      • 由多层的编码器(Encoder)组成（Transformer 的论文中使用了 6 层编码器，这里的层数 6 并不是固定的）
        • 每一个编码器，可以分为 2 层
          • Self-Attention Layer
          • Feed Forward Neural Network（前馈神经网络，缩写为 FFNN）
    • 解码部分(decoding component)
      • 由多层的解码器(Encoder)组成（Transformer 的论文中使用了 6 层解码器，这里的层数 6 并不是固定的）
        • 每一个解码器，可以分为 3层
          • Self-Attention Layer
          • Encoder-Decoder Attention 层，这个层能帮助解码器聚焦于输入句子的相关部分（类似于 seq2seq 模型 中的 Attention）。
          • Feed Forward Neural Network（前馈神经网络，缩写为 FFNN）

二、从细节来理解 Transformer

  • Transformer 的一个重要特性：
    • 每个位置的词向量经过编码器都有自己单独的路径。因此这些词向量在经过 Feed Forward 层中可以并行计算。

      三、 Self-Attention 整体理解

  1. # Standard PyTorch imports
  2. import numpy as np
  3. import torch
  4. import torch.nn as nn
  5. import torch.nn.functional as F
  6. import math, copy, time
  7. from torch.autograd import Variable
  9. # For plots
  10. %matplotlib inline
  11. import matplotlib.pyplot as plt
  13. ############################Model Architecture
  14. class EncoderDecoder(nn.Module):
  15.     """
  16.     A standard Encoder-Decoder architecture. Base for this and many 
  17.     other models.
  18.     """
  19.     def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
  20.         super(EncoderDecoder, self).__init__()
  21.         self.encoder = encoder
  22.         self.decoder = decoder
  23.         self.src_embed = src_embed
  24.         self.tgt_embed = tgt_embed
  25.         self.generator = generator
  27.     def forward(self, src, tgt, src_mask, tgt_mask):
  28.         "Take in and process masked src and target sequences."
  29.         return self.decode(self.encode(src, src_mask), src_mask,
  30.                             tgt, tgt_mask)
  32.     def encode(self, src, src_mask):
  33.         return self.encoder(self.src_embed(src), src_mask)
  35.     def decode(self, memory, src_mask, tgt, tgt_mask):
  36.         return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
  38. class Generator(nn.Module):
  39.     "Define standard linear + softmax generation step."
  40.     def __init__(self, d_model, vocab):
  41.         super(Generator, self).__init__()
  42.         self.proj = nn.Linear(d_model, vocab)
  44.     def forward(self, x):
  45.         return F.log_softmax(self.proj(x), dim=-1)
  47. ##############################Encoder and Decoder Stacks
  48. ##############################Encoder
  49. def clones(module, N):
  50.     "Produce N identical layers."
  51.     return nn.ModuleList([copy.deepcopy(module) for _ in range(N)])
  53. class Encoder(nn.Module):
  54.     "Core encoder is a stack of N layers"
  55.     def __init__(self, layer, N):
  56.         super(Encoder, self).__init__()
  57.         self.layers = clones(layer, N)
  58.         self.norm = LayerNorm(layer.size)
  60.     def forward(self, x, mask):
  61.         "Pass the input (and mask) through each layer in turn."
  62.         for layer in self.layers:
  63.             x = layer(x, mask)
  64.         return self.norm(x)
  66. class LayerNorm(nn.Module):
  67.     "Construct a layernorm module (See citation for details)."
  68.     def __init__(self, features, eps=1e-6):
  69.         super(LayerNorm, self).__init__()
  70.         self.a_2 = nn.Parameter(torch.ones(features))
  71.         self.b_2 = nn.Parameter(torch.zeros(features))
  72.         self.eps = eps
  74.     def forward(self, x):
  75.         mean = x.mean(-1, keepdim=True)
  76.         std = x.std(-1, keepdim=True)
  77.         return self.a_2 * (x - mean) / (std + self.eps) + self.b_2
  79. class SublayerConnection(nn.Module):
  80.     """
  81.     A residual connection followed by a layer norm.
  82.     Note for code simplicity the norm is first as opposed to last.
  83.     """
  84.     def __init__(self, size, dropout):
  85.         super(SublayerConnection, self).__init__()
  86.         self.norm = LayerNorm(size)
  87.         self.dropout = nn.Dropout(dropout)
  89.     def forward(self, x, sublayer):
  90.         "Apply residual connection to any sublayer with the same size."
  91.         return x + self.dropout(sublayer(self.norm(x)))
  93. class EncoderLayer(nn.Module):
  94.     "Encoder is made up of self-attn and feed forward (defined below)"
  95.     def __init__(self, size, self_attn, feed_forward, dropout):
  96.         super(EncoderLayer, self).__init__()
  97.         self.self_attn = self_attn
  98.         self.feed_forward = feed_forward
  99.         self.sublayer = clones(SublayerConnection(size, dropout), 2)
  100.         self.size = size
  102.     def forward(self, x, mask):
  103.         "Follow Figure 1 (left) for connections."
  104.         x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask))
  105.         return self.sublayer[1](x, self.feed_forward)
  108. ##############################Decoder
  109. class Decoder(nn.Module):
  110.     "Generic N layer decoder with masking."
  111.     def __init__(self, layer, N):
  112.         super(Decoder, self).__init__()
  113.         self.layers = clones(layer, N)
  114.         self.norm = LayerNorm(layer.size)
  116.     def forward(self, x, memory, src_mask, tgt_mask):
  117.         for layer in self.layers:
  118.             x = layer(x, memory, src_mask, tgt_mask)
  119.         return self.norm(x)
  121. class DecoderLayer(nn.Module):
  122.     "Decoder is made of self-attn, src-attn, and feed forward (defined below)"
  123.     def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
  124.         super(DecoderLayer, self).__init__()
  125.         self.size = size
  126.         self.self_attn = self_attn
  127.         self.src_attn = src_attn
  128.         self.feed_forward = feed_forward
  129.         self.sublayer = clones(SublayerConnection(size, dropout), 3)
  131.     def forward(self, x, memory, src_mask, tgt_mask):
  132.         "Follow Figure 1 (right) for connections."
  133.         m = memory
  134.         x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
  135.         x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
  136.         return self.sublayer[2](x, self.feed_forward)
  138. def subsequent_mask(size):
  139.     "Mask out subsequent positions."
  140.     attn_shape = (1, size, size)
  141.     subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')  # 返回函数的上三角矩阵
  142.     return torch.from_numpy(subsequent_mask) == 0
  145. #################################Attention
  146. #####################Scaled Dot-Product Attention
  147. def attention(query, key, value, mask=None, dropout=None):
  148.     "Compute 'Scaled Dot Product Attention'"
  149.     d_k = query.size(-1)
  150.     scores = torch.matmul(query, key.transpose(-2, -1)) \
  151.              / math.sqrt(d_k)
  152.     if mask is not None:
  153.         scores = scores.masked_fill(mask == 0, -1e9)
  154.     p_attn = F.softmax(scores, dim = -1)
  155.     if dropout is not None:
  156.         p_attn = dropout(p_attn)
  157.     return torch.matmul(p_attn, value), p_attn
  159. ###################Multi-head attention
  160. class MultiHeadedAttention(nn.Module):
  161.     def __init__(self, h, d_model, dropout=0.1):
  162.         "Take in model size and number of heads."
  163.         super(MultiHeadedAttention, self).__init__()
  164.         assert d_model % h == 0
  165.         # We assume d_v always equals d_k
  166.         self.d_k = d_model // h
  167.         self.h = h
  168.         self.linears = clones(nn.Linear(d_model, d_model), 4)
  169.         self.attn = None
  170.         self.dropout = nn.Dropout(p=dropout)
  172.     def forward(self, query, key, value, mask=None):
  173.         "Implements Figure 2"
  174.         if mask is not None:
  175.             # Same mask applied to all h heads.
  176.             mask = mask.unsqueeze(1)
  177.         nbatches = query.size(0)
  179.         # 1) 这一步qkv变化:[batch, L, d_model] ->[batch, h, L, d_model/h]
  180.         query, key, value = \
  181.             [l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
  182.              for l, x in zip(self.linears, (query, key, value))]
  184.         # 2) 计算注意力attn 得到attn*v 与attn 
  185.         x, self.attn = attention(query, key, value, mask=mask, 
  186.                                  dropout=self.dropout)
  188.         # 3) 上一步的结果合并在一起还原成原始输入序列的形状
  189.         x = x.transpose(1, 2).contiguous() \
  190.              .view(nbatches, -1, self.h * self.d_k)
  191.         # 最后再过一个线性层
  192.         return self.linears[-1](x)
  194. #######################Position-wise Feed-Forward Networks
  195. class PositionwiseFeedForward(nn.Module):
  196.     "Implements FFN equation."
  197.     def __init__(self, d_model, d_ff, dropout=0.1):
  198.         super(PositionwiseFeedForward, self).__init__()
  199.         self.w_1 = nn.Linear(d_model, d_ff)
  200.         self.w_2 = nn.Linear(d_ff, d_model)
  201.         self.dropout = nn.Dropout(dropout)
  203.     def forward(self, x):
  204.         return self.w_2(self.dropout(F.relu(self.w_1(x))))
  206. ######################Embeddings and Softmax
  207. class Embeddings(nn.Module):
  208.     def __init__(self, d_model, vocab):
  209.         super(Embeddings, self).__init__()
  210.         self.lut = nn.Embedding(vocab, d_model)
  211.         self.d_model = d_model
  213.     def forward(self, x):
  214.         return self.lut(x) * math.sqrt(self.d_model)
  217. ########################Positional Encoding
  218. class PositionalEncoding(nn.Module):
  219.     "Implement the PE function."
  220.     def __init__(self, d_model, dropout, max_len=5000):
  221.         super(PositionalEncoding, self).__init__()
  222.         self.dropout = nn.Dropout(p=dropout)
  224.         # Compute the positional encodings once in log space.
  225.         pe = torch.zeros(max_len, d_model)
  226.         position = torch.arange(0, max_len).unsqueeze(1)
  227.         div_term = torch.exp(torch.arange(0, d_model, 2) *
  228.                              -(math.log(10000.0) / d_model))
  229.         pe[:, 0::2] = torch.sin(position * div_term)
  230.         pe[:, 1::2] = torch.cos(position * div_term)
  231.         pe = pe.unsqueeze(0)
  232.         self.register_buffer('pe', pe)
  234.     def forward(self, x):
  235.         x = x + Variable(self.pe[:, :x.size(1)], 
  236.                          requires_grad=False)
  237.         return self.dropout(x)
  240. ################################Full Model
  241. def make_model(src_vocab, tgt_vocab, N=6, 
  242.                d_model=512, d_ff=2048, h=8, dropout=0.1):
  243.     "Helper: Construct a model from hyperparameters."
  244.     c = copy.deepcopy
  245.     attn = MultiHeadedAttention(h, d_model)
  246.     ff = PositionwiseFeedForward(d_model, d_ff, dropout)
  247.     position = PositionalEncoding(d_model, dropout)
  248.     model = EncoderDecoder(
  249.         Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout), N),
  250.         Decoder(DecoderLayer(d_model, c(attn), c(attn), 
  251.                              c(ff), dropout), N),
  252.         nn.Sequential(Embeddings(d_model, src_vocab), c(position)),
  253.         nn.Sequential(Embeddings(d_model, tgt_vocab), c(position)),
  254.         Generator(d_model, tgt_vocab))
  256.     # This was important from their code. 
  257.     # Initialize parameters with Glorot / fan_avg.
  258.     for p in model.parameters():
  259.         if p.dim() > 1:
  260.             nn.init.xavier_uniform(p)
  261.     return model
  263. tmp_model = make_model(10, 10, 2)
  264. tmp_model