透过机器翻译理解Transformer(四): 打造 Tran

2020-11-12 本文已影响0人博士伦2014

编者按：年初疫情在家期间开始大量阅读NLP领域的经典论文，在学习《Attention Is All You Need》时发现了一位现居日本的台湾数据科学家LeeMeng写的Transformer详解博客，理论讲解+代码实操+动画演示的写作风格，在众多文章中独树一帜，实为新手学习Transformer的上乘资料，在通读以及实操多遍之后，现在将其编辑整理成简体中文分享给大家。由于原文实在太长，为了便于阅读学习，这里将其分为四个部分：

在涉及代码部分，强烈推荐大家在Google的Colab Notebooks中实际操作一遍，之所以推荐Colab Notebooks是因为1).这里有免费可以使用的GPU资源；2). 可以避免很多安装包出错的问题

本节目录

1. 打造 Transformer：叠叠乐时间
  - 6.1 Position-wise Feed-Forward Networks
  - 6.2 Encoder layer：Encoder 小弟
  - 6.3 Decoder layer：Decoder 小弟
  - 6.4 Positional encoding：神奇数字
  - 6.5 Encoder
  - 6.6 Decoder
  - 6.7 第一个 Transformer
1. 定义损失函数与指标
1. 设置超参数
1. 设置 Optimizer
1. 实际训练以及定时存档
1. 实际进行英翻中
1. 可视化注意权重
1. 在你离开之前

6. 打造 Transformer：叠叠乐时间

以前我们曾提到深度学习模型就是一层层的几何运算过程。 Transformer 也不例外，刚才实现的 mutli-head attention layer 就是一个最明显的例子。而它正好是 Transformer 里头最重要的一层运算。

在这节我们会把 Transformer 里头除了注意力机制的其他运算通通实现成一个个的 layers，并将它们全部「叠」起来。

你可以通过下方的影片来了解接下来的实现顺序：

steps-to-build-transformer.gif

影片中左侧就是我们接下来会依序实现的 layers。 Transformer 是一种使用自注意力机制的 Seq2Seq 模型，里头包含了两个重要角色，分别为 Encoder 与 Decoder：

最初输入的英文序列会通过 Encoder 中 N 个 Encoder layers 并被转换成一个相同长度的序列。每个 layer 都会为自己的输入序列里头的子词产生新的 repr.，然后交给下一个 layer。
Decoder 在生成（预测）下一个中文子词时会一边观察 Encoder 输出序列里所有英文子词的 repr.，一边观察自己前面已经生成的中文子词。

值得一提的是，N = 1 （Encoder / Decoder layer 数目 = 1）时就是最阳春版的 Transformer。但在深度学习领域里头我们常常想对原始数据做多层的转换，因此会将 N 设为影片最后出现的 2 层或是 Transformer 论文中的 6 层 Encoder / Decoder layers。

Encoder 里头的 Encoder layer 里又分两个 sub-layers，而 Decoder 底下的 Decoder layer 则包含 3 个 sub-layers。真的是 layer layer 相扣。将这些 layers 的阶层关系简单列出来大概就长这样（位置 Encoding 等在实现时会做解释）：

Transformer
- Encoder
  - 输入 Embedding
  - 位置 Encoding
  - N 个 Encoder layers
    - sub-layer 1: Encoder 自注意力机制
    - sub-layer 2: Feed Forward
- Decoder
  - 输出 Embedding
  - 位置 Encoding
  - N 个 Decoder layers
    - sub-layer 1: Decoder 自注意力机制
    - sub-layer 2: Decoder-Encoder 注意力机制
    - sub-layer 3: Feed Forward
- Final Dense Layer

不过就像影片中显示的一样，实现的时候我们倾向从下往上叠上去。毕竟地基打得好，楼才盖得高，对吧？

6.1 Position-wise Feed-Forward Networks

如同影片中所看到的， Encoder layer 跟 Decoder layer 里头都各自有一个 Feed Forward 的元件。此元件构造简单，不用像前面的multi-head attention 建立定制化的keras layer，只需要写一个Python 函数让它在被调用的时候返回一个新的tf.keras.Sequential 模型给我们即可：

# 建立 Transformer 里 Encoder / Decoder layer 都有使用到的 Feed Forward 元件
def point_wise_feed_forward_network(d_model, dff):
  
  # 此 FFN 对输入做两个线性转换，中间加了一个 ReLU activation func
  return tf.keras.Sequential([
      tf.keras.layers.Dense(dff, activation='relu'),  # (batch_size, seq_len, dff)
      tf.keras.layers.Dense(d_model)  # (batch_size, seq_len, d_model)
  ])

此函数在每次被调用的时候都会返回一组新的全连接前馈神经网路（Fully-connected Feed Forward Network，FFN），其输入张量与输出张量的最后一个维度皆为d_model ，而在FFN 中间层的维度则为dff。一般会让 dff大于 d_model，让 FFN 从输入的d_model维度里头撷取些有用的信息。在论文中d_model为 512，dff 则为 4 倍的d_model: 2048。两个都是可以调整的超参数。

让我们建立一个 FFN 试试：

batch_size = 64
seq_len = 10
d_model = 512
dff = 2048

x = tf.random.uniform((batch_size, seq_len, d_model))
ffn = point_wise_feed_forward_network(d_model, dff)
out = ffn(x)
print("x.shape:", x.shape)
print("out.shape:", out.shape)

x.shape: (64, 10, 512)
out.shape: (64, 10, 512)

在输入张量的最后一维已经是 d_model 的情况，FFN 的输出张量基本上会跟输入一模一样：

输入：（batch_size, seq_len, d_model）
输出：（batch_size, seq_len, d_model）

FFN 输出 / 输入张量的 shape 相同很容易理解。比较没那么明显的是这个 FFN 事实上对序列中的所有位置做的线性转换都是一样的。我们可以假想一个 2 维的 duumy_sentence，里头有 5 个以 4 维向量表示的子词：

d_model = 4 # FFN 的输入输出张量的最后一维皆为 `d_model`
dff = 6

# 建立一个小 FFN
small_ffn = point_wise_feed_forward_network(d_model, dff)
# 懂子词梗的站出来
dummy_sentence = tf.constant([[5, 5, 6, 6], 
                              [5, 5, 6, 6], 
                              [9, 5, 2, 7], 
                              [9, 5, 2, 7],
                              [9, 5, 2, 7]], dtype=tf.float32)
small_ffn(dummy_sentence)

<tf.Tensor: shape=(5, 4), dtype=float32, numpy=
array([[ 2.8674245, -2.174698 , -1.3073452, -6.4233937],
       [ 2.8674245, -2.174698 , -1.3073452, -6.4233937],
       [ 3.650207 , -0.973258 , -2.4126565, -6.5094995],
       [ 3.650207 , -0.973258 , -2.4126565, -6.5094995],
       [ 3.650207 , -0.973258 , -2.4126565, -6.5094995]], dtype=float32)>

你会发现同一个子词不会因为位置的改变而造成 FFN 的输出结果产生差异。但因为我们实际上会有多个 Encoder / Decoder layers，而每个 layers 都会有不同参数的 FFN，因此每个 layer 里头的 FFN 做的转换都会有所不同。

值得一提的是，尽管对所有位置的子词都做一样的转换，但是这个转换是独立进行的，因此被称作 Position-wise Feed-Forward Networks。

6.2 Encoder layer：Encoder 小弟

有了 Multi-Head Attention（MHA）以及 Feed-Forward Network（FFN），我们事实上已经可以实现第一个 Encoder layer 了。让我们复习一下这 layer 里头有什么重要元件：

Encoder layer 里的重要元件

我想上面的动画已经很清楚了。一个 Encoder layer 里头会有两个 sub-layers，分别为 MHA 以及 FFN。在 Add & Norm 步骤里头，每个 sub-layer 会有一个残差连结（residual connection）来帮助减缓梯度消失（Gradient Vanishing）的问题。接着两个 sub-layers 都会针对最后一维 d_model 做 layer normalization，将 batch 里头每个子词的输出独立做转换，使其平均与标准差分别靠近 0 和 1 之后输出。

另外在将 sub-layer 的输出与其输入相加之前，我们还会做点 regularization，对该 sub-layer 的输出使用 dropout。

总结一下。如果输入是 x，最后输出写作out的话，则每个 sub-layer 的处理逻辑如下：

sub_layer_out = Sublayer(x)
sub_layer_out = Dropout(sub_layer_out)
out = LayerNorm(x + sub_layer_out)

Sublayer 则可以是 MHA 或是 FFN。现在让我们看看 Encoder layer 的实现：

# Encoder 里头会有 N 个 EncoderLayers，而每个 EncoderLayer 里又有两个 sub-layers: MHA & FFN
class EncoderLayer(tf.keras.layers.Layer):
    
    # Transformer 论文内预设 dropout rate 为 0.1
    def __init__(self, d_model, num_heads, dff, rate=0.1):
        super(EncoderLayer, self).__init__()

        self.mha = MultiHeadAttention(d_model, num_heads)
        self.ffn = point_wise_feed_forward_network(d_model, dff)

       # layer norm 很常在 RNN-based 的模型被使用。一个 sub-layer 一个 layer norm
        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
        self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    
        # 一样，一个 sub-layer 一个 dropout layer
        self.dropout1 = tf.keras.layers.Dropout(rate)
        self.dropout2 = tf.keras.layers.Dropout(rate)

    
    # 需要丢入 `training` 参数是因为 dropout 在训练以及测试的行为有所不同
    def call(self, x, training, mask):
        # 除了 `attn`，其他张量的 shape 皆为 (batch_size, input_seq_len, d_model)
        # attn.shape == (batch_size, num_heads, input_seq_len, input_seq_len)
    
        # sub-layer 1: MHA
        # Encoder 利用注意机制关注自己当前的序列，因此 v, k, q 全部都是自己
        # 另外别忘了我们还需要 padding mask 来遮住输入序列中的 <pad> token
        attn_output, attn = self.mha(x, x, x, mask)  
        attn_output = self.dropout1(attn_output, training=training) 
        out1 = self.layernorm1(x + attn_output)  
    
        # sub-layer 2: FFN
        ffn_output = self.ffn(out1) 
        ffn_output = self.dropout2(ffn_output, training=training)  # 記得 training
        out2 = self.layernorm2(out1 + ffn_output)
    
        return out2

跟当初 MHA layer 的实作比起来轻松多了，对吧？

基本上 Encoder layer 里头就是两个架构一模一样的 sub-layer，只差一个是 MHA，一个是 FFN。另外为了方便 residual connection 的计算，所有 sub-layers 的输出维度都是 d_model。而 sub-layer 内部产生的维度当然就随我们开心啦！我们可以为 FFN 设置不同的 dff 值，也能设定不同的 num_heads 来改变 MHA 内部每个 head 里头的维度。

论文里头的 d_model为 512，而我们 demo 用的英文词嵌入张量的d_model 维度则为 4：

# 之后可以调的超参数。这边为了 demo 设小一点
d_model = 4
num_heads = 2
dff = 8

# 新建一个使用上述参数的 Encoder Layer
enc_layer = EncoderLayer(d_model, num_heads, dff)
padding_mask = create_padding_mask(inp)  # 建立一个当前输入 batch 使用的 padding mask
enc_out = enc_layer(emb_inp, training=False, mask=padding_mask)  # (batch_size, seq_len, d_model)

print("inp:", inp)
print("-" * 20)
print("padding_mask:", padding_mask)
print("-" * 20)
print("emb_inp:", emb_inp)
print("-" * 20)
print("enc_out:", enc_out)
assert emb_inp.shape == enc_out.shape

inp: tf.Tensor(
[[8113  103    9 1066 7903 8114    0    0]
 [8113   16 4111 6735   12 2750 7903 8114]], shape=(2, 8), dtype=int64)
 --------------------
padding_mask: tf.Tensor(
[[[[0. 0. 0. 0. 0. 0. 1. 1.]]]
[[[0. 0. 0. 0. 0. 0. 0. 0.]]]], shape=(2, 1, 1, 8), dtype=float32)
--------------------
emb_inp: tf.Tensor(
[[[ 0.0041508   0.04106052  0.00270988 -0.00628465]
  [ 0.0261193   0.04892724 -0.03637441  0.00032102]
  [-0.0315491   0.03012072 -0.03764988 -0.00832593]
  [-0.00863073  0.01537497  0.00647591  0.01622475]
  [ 0.01064278  0.02867876  0.0471475   0.02418466]
  [-0.0357633  -0.02500458  0.00584758  0.00984917]
  [ 0.02766568 -0.02055204  0.0366873  -0.04519999]
  [ 0.02766568 -0.02055204  0.0366873  -0.04519999] 
  [ 0.0041508   0.04106052  0.00270988 -0.00628465]
  [-0.03440493  0.0245572  -0.04154334  0.01249687]
  [-0.04102417 -0.04214551 -0.03087332  0.03536062]
  [ 0.00288613 -0.00550915  0.02198391 -0.02721313]
  [ 0.03594044 -0.02207484  0.00774273 -0.01938369]
  [-0.00556026  0.04242435  0.03270287 -0.00513189]
  [ 0.01064278  0.02867876  0.0471475   0.02418466]
  [-0.0357633  -0.02500458  0.00584758  0.00984917]]], shape=(2, 8, 4), dtype=float32)
 --------------------
enc_out: tf.Tensor(
[[[-0.1656846   1.4814154  -1.3332843   0.01755357]
  [ 0.05347645  1.2417278  -1.5466218   0.25141746]
  [-0.8423737   1.4621214  -1.0028969   0.3831491 ]
  [-1.1612244   0.4753281  -0.7035671   1.3894634 ]
  [-1.0288012  -0.7085241   1.5507177   0.1866076 ]
  [-0.5757953  -1.1105288   0.13135773  1.5549664 ]
  [ 1.5314106  -0.519994    0.1549343  -1.1663508 ]
  [ 1.5314106  -0.519994    0.1549343  -1.1663508 ]]

 [[-0.34800935  1.5336158  -1.234706    0.04909949]
  [-0.97635764  1.3417366  -0.9507113   0.58533245]
  [-0.53843904 -0.48348504 -0.7043885   1.7263125 ]
  [ 1.208463   -0.2577439   0.529937   -1.4806561 ]
  [ 1.6743237  -0.9758253  -0.33426592 -0.36423233]
  [-1.0195854   1.6443692  -0.13730906 -0.48747474]
  [-1.4697037  -0.00313468  1.3509609   0.12187762]
  [-0.8544105  -0.8589976   0.12724805  1.5861602 ]]], shape=(2, 8, 4), dtype=float32)

在本来的输入维度即为 d_model 的情况下，Encoder layer 就是给我们一个一模一样 shape 的张量。当然，实际上内部透过 MHA 以及 FFN sub-layer 的转换，每个子词的 repr. 都大幅改变了。

有了 Encoder layer，接着让我们看看 Decoder layer 的实现。

6.3 Decoder layer：Decoder 小弟

一个 Decoder layer 里头有 3 个 sub-layers：

Decoder layer 自身的 Masked MHA 1
Decoder layer 关注 Encoder 输出序列的 MHA 2
FFN

你也可以看一下影片来回顾它们所在的位置：

Decoder layer 中的 sub-layers

跟实现 Encoder layer 时一样，每个 sub-layer 的逻辑同下：

sub_layer_out = Sublayer(x)
sub_layer_out = Dropout(sub_layer_out)
out = LayerNorm(x + sub_layer_out)

Decoder layer 用 MHA 1 来关注输出序列，查询 Q、键值 K 以及值 V 都是自己。而之所以有个 masked 是因为（中文）输出序列除了跟（英文）输入序列一样需要 padding mask 以外，还需要 look ahead mask 来避免 Decoder layer 关注到未来的子词。 look ahead mask 在前面章节已经有详细说明了。

MHA1 处理完的输出序列会成为 MHA 2 的 Q，而 K 与 V 则使用 Encoder 的输出序列。这个运算的概念是让一个 Decoder layer 在生成新的中文子词时先参考先前已经产生的中文子词，并为当下要生成的子词产生一个包含前文语义的 repr. 。接着将此 repr. 拿去跟 Encoder 那边的英文序列做匹配，看当下子词的 repr. 有多好并予以修正。

用简单点的说法就是 Decoder 在生成中文子词时除了参考自己已经生成的中文子词以外，也会去关注 Encoder 输出的英文子词（的 repr.）。

# Decoder 里头会有 N 个 DecoderLayer，
# 而 DecoderLayer 又有三个 sub-layers: 自注意的 MHA, 关注 Encoder 输出的 MHA & FFN

class DecoderLayer(tf.keras.layers.Layer):
  def __init__(self, d_model, num_heads, dff, rate=0.1):
    super(DecoderLayer, self).__init__()

    # 3 個 sub-layers 的主角們
    self.mha1 = MultiHeadAttention(d_model, num_heads)
    self.mha2 = MultiHeadAttention(d_model, num_heads)
    self.ffn = point_wise_feed_forward_network(d_model, dff)
 
    # 定義每個 sub-layer 用的 LayerNorm
    self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm2 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    self.layernorm3 = tf.keras.layers.LayerNormalization(epsilon=1e-6)
    
    # 定義每個 sub-layer 用的 Dropout
    self.dropout1 = tf.keras.layers.Dropout(rate)
    self.dropout2 = tf.keras.layers.Dropout(rate)
    self.dropout3 = tf.keras.layers.Dropout(rate)
    
    
  def call(self, x, enc_output, training, 
           combined_mask, inp_padding_mask):
    # 所有 sub-layers 的主要輸出皆為 (batch_size, target_seq_len, d_model)
    # enc_output 為 Encoder 輸出序列，shape 為 (batch_size, input_seq_len, d_model)
    # attn_weights_block_1 則為 (batch_size, num_heads, target_seq_len, target_seq_len)
    # attn_weights_block_2 則為 (batch_size, num_heads, target_seq_len, input_seq_len)

    # sub-layer 1: Decoder layer 自己對輸出序列做注意力。
    # 我們同時需要 look ahead mask 以及輸出序列的 padding mask 
    # 來避免前面已生成的子詞關注到未來的子詞以及 <pad>
    attn1, attn_weights_block1 = self.mha1(x, x, x, combined_mask)
    attn1 = self.dropout1(attn1, training=training)
    out1 = self.layernorm1(attn1 + x)
    
    # sub-layer 2: Decoder layer 關注 Encoder 的最後輸出
    # 記得我們一樣需要對 Encoder 的輸出套用 padding mask 避免關注到 <pad>
    attn2, attn_weights_block2 = self.mha2(
        enc_output, enc_output, out1, inp_padding_mask)  # (batch_size, target_seq_len, d_model)
    attn2 = self.dropout2(attn2, training=training)
    out2 = self.layernorm2(attn2 + out1)  # (batch_size, target_seq_len, d_model)
    
    # sub-layer 3: FFN 部分跟 Encoder layer 完全一樣
    ffn_output = self.ffn(out2)  # (batch_size, target_seq_len, d_model)

    ffn_output = self.dropout3(ffn_output, training=training)
    out3 = self.layernorm3(ffn_output + out2)  # (batch_size, target_seq_len, d_model)
    
    # 除了主要輸出 `out3` 以外，輸出 multi-head 注意權重方便之後理解模型內部狀況
    return out3, attn_weights_block1, attn_weights_block2

Decoder layer 的实现跟 Encoder layer 大同小异，不过还是有几点细节特别需要注意：

在做 Masked MHA（MHA 1）的时候我们需要同时套用两种遮罩：输出序列的 padding mask 以及 look ahead mask。因此 Decoder layer 预期的遮罩是两者结合的combined_mask
MHA 1 因为是 Decoder layer 关注自己，multi-head attention 的参数 v、k 以及 q都是 x
MHA 2 是 Decoder layer 关注 Encoder 输出序列，因此，multi-head attention 的参数 v、k 为 enc_output，q 则为 MHA 1 sub-layer 的结果 out1

产生comined_mask也很简单，我们只要把两个遮罩取大的即可：

tar_padding_mask = create_padding_mask(tar)
look_ahead_mask = create_look_ahead_mask(tar.shape[-1])
combined_mask = tf.maximum(tar_padding_mask, look_ahead_mask)

print("tar:", tar)
print("-" * 20)
print("tar_padding_mask:", tar_padding_mask)
print("-" * 20)
print("look_ahead_mask:", look_ahead_mask)
print("-" * 20)
print("combined_mask:", combined_mask)

tar: tf.Tensor(
[[4205   10  241   86   27    3 4206    0    0    0]
 [4205  165  489  398  191   14    7  560    3 4206]], shape=(2, 10), dtype=int64)
--------------------
tar_padding_mask: tf.Tensor(
[[[[0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]]]
[[[0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]]], shape=(2, 1, 1, 10), dtype=float32)
--------------------
look_ahead_mask: tf.Tensor(
 [[0. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
  [0. 0. 1. 1. 1. 1. 1. 1. 1. 1.]
  [0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
  [0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
  [0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
  [0. 0. 0. 0. 0. 0. 1. 1. 1. 1.]
  [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
  [0. 0. 0. 0. 0. 0. 0. 0. 1. 1.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]], shape=(10, 10), dtype=float32)
 --------------------
combined_mask: tf.Tensor(
[[[[0. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
   [0. 0. 1. 1. 1. 1. 1. 1. 1. 1.]
   [0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
   [0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
   [0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 1. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]]]

[[[0. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
  [0. 0. 1. 1. 1. 1. 1. 1. 1. 1.]
  [0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
  [0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
  [0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
  [0. 0. 0. 0. 0. 0. 1. 1. 1. 1.]
  [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
  [0. 0. 0. 0. 0. 0. 0. 0. 1. 1.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]
  [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]]], shape=(2, 1, 10, 10), dtype=float32)

注意 combined_mask 的 shape 以及里头遮罩所在的位置。利用 broadcasting 我们将 combined_mask 的 shape 也扩充到 4 维：

(batch_size, num_heads, seq_len_tar, seq_len_tar)= (2, 1, 10, 10)

这方便之后 multi-head attention 的计算。另外因为我们 demo 的中文 batch 里头的第一个句子有 <pad>，combined_mask除了 look ahead 的效果以外还加了 padding mask。

因为刚刚实现的是 Decoder layer，这次让我们把中文（目标语言）的词嵌入张量以及相关的遮罩丢进去看看：

# 超参数
d_model = 4
num_heads = 2
dff = 8
dec_layer = DecoderLayer(d_model, num_heads, dff)

# 来源、目标语言的序列都需要 padding mask
inp_padding_mask = create_padding_mask(inp)
tar_padding_mask = create_padding_mask(tar)

# masked MHA 用的遮罩，把 padding 跟未来子词都盖住
look_ahead_mask = create_look_ahead_mask(tar.shape[-1])
combined_mask = tf.maximum(tar_padding_mask, look_ahead_mask)

# 实际初始一个 decoder layer 并做 3 个 sub-layers 的计算
dec_out, dec_self_attn_weights, dec_enc_attn_weights = dec_layer(
    emb_tar, enc_out, False, combined_mask, inp_padding_mask)

print("emb_tar:", emb_tar)
print("-" * 20)
print("enc_out:", enc_out)
print("-" * 20)
print("dec_out:", dec_out)
assert emb_tar.shape == dec_out.shape
print("-" * 20)
print("dec_self_attn_weights.shape:", dec_self_attn_weights.shape)
print("dec_enc_attn_weights:", dec_enc_attn_weights.shape)

emb_tar: tf.Tensor(
[[[-0.00084939 -0.02029408 -0.04978932 -0.02889797]
  [-0.01320463  0.00070287  0.00797179 -0.00549082]
  [-0.01859868 -0.04142375  0.02479618 -0.00794141]
  [ 0.04030085 -0.04564189 -0.03584541 -0.04098076]
  [ 0.02629851  0.01072141 -0.01055797  0.04544314]
  [-0.00223017  0.02058548  0.01649131 -0.01385387]
  [ 0.00302396 -0.03152249  0.0396189  -0.03036447]
  [ 0.00433234  0.04481849  0.04129448  0.04720709]
  [ 0.00433234  0.04481849  0.04129448  0.04720709]
  [ 0.00433234  0.04481849  0.04129448  0.04720709]]

 [[-0.00084939 -0.02029408 -0.04978932 -0.02889797]
  [-0.04702241  0.01816512 -0.02416607 -0.01993601]
  [ 0.04391925 -0.03093947 -0.01225864 -0.03517971]
  [ 0.03755457  0.00626134  0.04324439  0.00490584]
  [ 0.00495391 -0.03399891  0.04144105  0.02539945]
  [ 0.0282723  -0.0164601  -0.00685417 -0.02280444]
  [ 0.04738505 -0.01041915 -0.02054645 -0.00066562]
  [-0.00438491  0.02117647 -0.04890387 -0.01620366]
  [-0.00223017  0.02058548  0.01649131 -0.01385387]
  [ 0.00302396 -0.03152249  0.0396189  -0.03036447]]], shape=(2, 10, 4), dtype=float32)
--------------------
enc_out: tf.Tensor(
[[[-0.1656846   1.4814154  -1.3332843   0.01755357]
  [ 0.05347645  1.2417278  -1.5466218   0.25141746]
  [-0.8423737   1.4621214  -1.0028969   0.3831491 ]
  [-1.1612244   0.4753281  -0.7035671   1.3894634 ]
  [-1.0288012  -0.7085241   1.5507177   0.1866076 ]
  [-0.5757953  -1.1105288   0.13135773  1.5549664 ]
  [ 1.5314106  -0.519994    0.1549343  -1.1663508 ]
  [ 1.5314106  -0.519994    0.1549343  -1.1663508 ]]

 [[-0.34800935  1.5336158  -1.234706    0.04909949]
  [-0.97635764  1.3417366  -0.9507113   0.58533245]
  [-0.53843904 -0.48348504 -0.7043885   1.7263125 ]
  [ 1.208463   -0.2577439   0.529937   -1.4806561 ]
  [ 1.6743237  -0.9758253  -0.33426592 -0.36423233]
  [-1.0195854   1.6443692  -0.13730906 -0.48747474]
  [-1.4697037  -0.00313468  1.3509609   0.12187762]
  [-0.8544105  -0.8589976   0.12724805  1.5861602 ]]], shape=(2, 8, 4), dtype=float32)
--------------------
dec_out: tf.Tensor(
[[[ 1.2991211   0.6467309  -0.99355525 -0.9522968 ]
  [-0.68756247 -0.44788587  1.7257465  -0.5902982 ]
  [ 0.21567897 -1.6887752   0.6456864   0.8274099 ]
  [ 1.3437784  -1.2335085  -0.6324715   0.52220154]
  [ 0.5747509  -1.1840664  -0.71563935  1.3249549 ]
  [-0.4092589   0.41854465  1.3476295  -1.3569155 ]
  [ 0.47711575 -1.7147235   0.8007993   0.43680844]
  [-1.132223   -0.82594645  1.222668    0.73550147]
  [-1.132223   -0.82594645  1.222668    0.73550147]
  [-1.132223   -0.82594645  1.222668    0.73550147]]

 [[ 1.3999227   0.49366176 -0.9038905  -0.989694  ]
  [-0.86972106  1.1954616   0.77558595 -1.1013266 ]
  [ 1.6006857  -1.068229   -0.5445589   0.01210219]
  [ 0.7155672  -1.6947896   0.750581    0.2286414 ]
  [ 0.1127052  -1.6265972   0.4442618   1.0696301 ]
  [ 1.4985088  -1.2589391  -0.38515666  0.14558706]
  [ 1.3210055  -0.90092945 -1.033153    0.6130771 ]
  [-0.0833452   1.6214814  -1.0698308  -0.4683055 ]
  [-0.4484089   0.17643274  1.5017867  -1.2298107 ]
  [ 0.44141728 -1.6816832   0.94259256  0.2976733 ]]], shape=(2, 10, 4), dtype=float32)
--------------------
dec_self_attn_weights.shape: (2, 2, 10, 10)
dec_enc_attn_weights: (2, 2, 10, 8)

跟 Encoder layer 相同，Decoder layer 输出张量的最后一维也是 d_model。而 dec_self_attn_weights 则代表着 Decoder layer 的自注意力权重，因此最后两个维度皆为中文序列的长度 10；而 dec_enc_attn_weights因为 Encoder 输出序列的长度为8，最后一维即为 8。

都读到这里了，判断每一维的物理意义对你来说应该是小菜一碟了。

6.4 Positional encoding：神奇数字

透过多层的自注意力层，Transformer 在处理序列时里头所有子词都是「天涯若比邻」：想要关注序列中任何位置的信息只要 O(1) 就能办到。这让 Transformer 能很好地 model 序列中长距离的依赖关系（long-range dependencise）。但反过来说 Transformer 则无法 model 序列中字词的顺序关系，所以我们得额外加入一些「位置信息」给 Transformer。

这个信息被称作位置编码（Positional Encoding），实作上是直接加到最一开始的英文 / 中文词嵌入向量（word embedding）里头。其直观的想法是想办法让被加入位置编码的 word embedding 在d_model维度的空间里头不只会因为语义相近而靠近，也会因为位置靠近而在该空间里头靠近。

论文里头使用的位置编码的公式如下：

position-encoding-equation.jpg

论文里头提到他们之所以这样设计位置编码（Positional Encoding, PE）是因为这个函数有个很好的特性：给定任一位置pos 的位置编码PE(pos)，跟它距离k个单位的位置pos + k 的位置编码PE(pos + k)可以表示为PE(pos) 的一个线性函数（linear function）。

因此透过在 word embedding 里加入这样的信息，作者们认为可以帮助 Transformer 学会 model 序列中的子词的相对位置关系。

就算我们无法自己想出论文里头的位置编码公式，还是可以直接把 TensorFlow 官方的实现搬过来使用：

# 以下直接参考 TensorFlow 官方 tutorial
def get_angles(pos, i, d_model):
  angle_rates = 1 / np.power(10000, (2 * (i//2)) / np.float32(d_model))
  return pos * angle_rates

def positional_encoding(position, d_model):
  angle_rads = get_angles(np.arange(position)[:, np.newaxis],
                          np.arange(d_model)[np.newaxis, :],
                          d_model)
  
  # apply sin to even indices in the array; 2i
  sines = np.sin(angle_rads[:, 0::2])
  
  # apply cos to odd indices in the array; 2i+1
  cosines = np.cos(angle_rads[:, 1::2])
  
  pos_encoding = np.concatenate([sines, cosines], axis=-1)
  
  pos_encoding = pos_encoding[np.newaxis, ...]
    
  return tf.cast(pos_encoding, dtype=tf.float32)


seq_len = 50
d_model = 512

pos_encoding = positional_encoding(seq_len, d_model)
pos_encoding

<tf.Tensor: shape=(1, 50, 512), dtype=float32, numpy=
array([[[ 0.        ,  0.        ,  0.        , ...,  1.        ,
          1.        ,  1.        ],
        [ 0.84147096,  0.8218562 ,  0.8019618 , ...,  1.        ,
          1.        ,  1.        ],
        [ 0.9092974 ,  0.9364147 ,  0.95814437, ...,  1.        ,
          1.        ,  1.        ],
        ...,
        [ 0.12357312,  0.97718984, -0.24295525, ...,  0.9999863 ,
          0.99998724,  0.99998814],
        [-0.76825464,  0.7312359 ,  0.63279754, ...,  0.9999857 ,
          0.9999867 ,  0.9999876 ],
        [-0.95375264, -0.14402692,  0.99899054, ...,  0.9999851 ,
          0.9999861 ,  0.9999871 ]]], dtype=float32)>

一路看下来你应该也可以猜到位置编码的每一维意义了：

第 1 维代表 batch_size，之后可以 broadcasting
第 2 维是序列长度，我们会为每个在输入 / 输出序列里头的子词都加入位置编码
第 3 维跟词嵌入向量同维度

因为是要跟词嵌入向量相加，位置编码的维度也得是 d_model。我们也可以把位置编码画出感受一下：

plt.pcolormesh(pos_encoding[0], cmap='RdBu')
plt.xlabel('d_model')
plt.xlim((0, 512))
plt.ylabel('Position')
plt.colorbar()
plt.show()

positional-encoding

这图你应该在很多教学文章以及教授的影片里都看过了。就跟我们前面看过的各种 2 维矩阵相同，x 轴代表着跟词嵌入向量相同的维度 d_model，y 轴则代表序列中的每个位置。之后我们会看输入 / 输出序列有多少个子词，就加入几个位置编码。

关于位置编码我们现在只需要知道这些就够了，但如果你想知道更多相关的数学计算，可以参考这个笔记本。

6.5 Encoder

Encoder 里头主要包含了 3 个元件：

输入的词嵌入层
位置编码
N 个 Encoder layers

大部分的工作都交给 Encoder layer 小弟做了，因此 Encoder 的实现很单纯：

class Encoder(tf.keras.layers.Layer):
  # Encoder 的初始參數除了本來就要給 EncoderLayer 的參數還多了：
  # - num_layers: 決定要有幾個 EncoderLayers, 前面影片中的 `N`
  # - input_vocab_size: 用來把索引轉成詞嵌入向量
  def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, 
               rate=0.1):
    super(Encoder, self).__init__()

    self.d_model = d_model
    
    self.embedding = tf.keras.layers.Embedding(input_vocab_size, d_model)
    self.pos_encoding = positional_encoding(input_vocab_size, self.d_model)
    
    # 建立 `num_layers` 個 EncoderLayers
    self.enc_layers = [EncoderLayer(d_model, num_heads, dff, rate) 
                       for _ in range(num_layers)]

    self.dropout = tf.keras.layers.Dropout(rate)
        
  def call(self, x, training, mask):
    # 輸入的 x.shape == (batch_size, input_seq_len)
    # 以下各 layer 的輸出皆為 (batch_size, input_seq_len, d_model)
    input_seq_len = tf.shape(x)[1]
    
    # 將 2 維的索引序列轉成 3 維的詞嵌入張量，並依照論文乘上 sqrt(d_model)
    # 再加上對應長度的位置編碼
    x = self.embedding(x)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :input_seq_len, :]

    # 對 embedding 跟位置編碼的總合做 regularization
    # 這在 Decoder 也會做
    x = self.dropout(x, training=training)
    
    # 通過 N 個 EncoderLayer 做編碼
    for i, enc_layer in enumerate(self.enc_layers):
      x = enc_layer(x, training, mask)
      # 以下只是用來 demo EncoderLayer outputs
      #print('-' * 20)
      #print(f"EncoderLayer {i + 1}'s output:", x)
      
    
    return x

比较值得注意的是我们依照论文将 word embedding 乘上 sqrt(d_model)，并在 embedding 跟位置编码相加以后通过 dropout 层来达到 regularization 的效果。

现在我们可以直接将索引序列 inp 丢入 Encoder：

# 超参数
num_layers = 2 # 2 層的 Encoder
d_model = 4
num_heads = 2
dff = 8
input_vocab_size = subword_encoder_en.vocab_size + 2 # 记得加上 <start>, <end>

# 初始化一个 Encoder
encoder = Encoder(num_layers, d_model, num_heads, dff, input_vocab_size)

# 将 2 维的索引序列丢入 Encoder 做编码
enc_out = encoder(inp, training=False, mask=None)
print("inp:", inp)
print("-" * 20)
print("enc_out:", enc_out)

inp: tf.Tensor(
[[8113  103    9 1066 7903 8114    0    0]
 [8113   16 4111 6735   12 2750 7903 8114]], shape=(2, 8), dtype=int64)
--------------------
enc_out: tf.Tensor(
[[[-0.7849332  -0.5919684  -0.33270505  1.7096066 ]
  [-0.5070654  -0.5110136  -0.7082318   1.726311  ]
  [-0.39270183 -0.03102639 -1.158362    1.5820901 ]
  [-0.5561629   0.38050282 -1.2407898   1.4164499 ]
  [-0.90432     0.19381052 -0.8472892   1.5577985 ]
  [-0.97321564 -0.22992788 -0.4652462   1.6683896 ]
  [-0.84681976 -0.5434473  -0.31013608  1.7004032 ]
  [-0.62432766 -0.56790507 -0.539001    1.7312336 ]]

 [[-0.77423775 -0.6076471  -0.32800597  1.7098908 ]
  [-0.47978252 -0.5615605  -0.68602914  1.7273722 ]
  [-0.30068305 -0.07366991 -1.1973959   1.5717487 ]
  [-0.5147841   0.2787246  -1.2290851   1.4651446 ]
  [-0.89634496  0.2675462  -0.8954112   1.52421   ]
  [-0.97553635 -0.22618684 -0.4656965   1.6674198 ]
  [-0.87600434 -0.5448401  -0.27099532  1.6918398 ]
  [-0.60130465 -0.5993665  -0.5306774   1.7313484 ]]], shape=(2, 8, 4), dtype=float32)

注意因为 Encoder 已经包含了词嵌入层，因此我们不用再像调用 Encoder layer 时一样还得自己先做 word embedding。现在的输入及输出张量为：

输入：（batch_size, seq_len）
输出：（batch_size, seq_len, d_model）

有了 Encoder，我们之后就可以直接把 2 维的索引序列 inp丢入 Encoder，让它帮我们把里头所有的英文序列做一连串的转换。

6.6 Decoder

Decoder layer 本来就只跟 Encoder layer 差在一个 MHA，而这逻辑被包起来以后调用它的 Decoder 做的事情就跟 Encoder 基本上没有两样了。

在 Decoder 里头我们只需要建立一个专门给中文用的词嵌入层以及位置编码即可。我们在调用每个 Decoder layer 的时候也顺便把其注意力权重存下来，方便我们了解模型训练完后是怎么做翻译的。

以下则是实现：

class Decoder(tf.keras.layers.Layer):
  # 初始參數跟 Encoder 只差在用 `target_vocab_size` 而非 `inp_vocab_size`
  def __init__(self, num_layers, d_model, num_heads, dff, target_vocab_size, 
               rate=0.1):
    super(Decoder, self).__init__()

    self.d_model = d_model
    
    # 為中文（目標語言）建立詞嵌入層
    self.embedding = tf.keras.layers.Embedding(target_vocab_size, d_model)
    self.pos_encoding = positional_encoding(target_vocab_size, self.d_model)
    
    self.dec_layers = [DecoderLayer(d_model, num_heads, dff, rate) 
                       for _ in range(num_layers)]
    self.dropout = tf.keras.layers.Dropout(rate)
  
  # 呼叫時的參數跟 DecoderLayer 一模一樣
  def call(self, x, enc_output, training, 
           combined_mask, inp_padding_mask):
    
    tar_seq_len = tf.shape(x)[1]
    attention_weights = {}  # 用來存放每個 Decoder layer 的注意權重
    
    # 這邊跟 Encoder 做的事情完全一樣
    x = self.embedding(x)  # (batch_size, tar_seq_len, d_model)
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x += self.pos_encoding[:, :tar_seq_len, :]
    x = self.dropout(x, training=training)

    
    for i, dec_layer in enumerate(self.dec_layers):
      x, block1, block2 = dec_layer(x, enc_output, training,
                                    combined_mask, inp_padding_mask)
      
      # 將從每個 Decoder layer 取得的注意權重全部存下來回傳，方便我們觀察
      attention_weights['decoder_layer{}_block1'.format(i + 1)] = block1
      attention_weights['decoder_layer{}_block2'.format(i + 1)] = block2
    
    # x.shape == (batch_size, tar_seq_len, d_model)
    return x, attention_weights

接着让我们初始并调用一个 Decoder 看看：

# 超参数
num_layers = 2 # 2 层的 Decoder
d_model = 4
num_heads = 2
dff = 8
target_vocab_size = subword_encoder_zh.vocab_size + 2 # 记得加上 <start>, <end>

# 遮罩
inp_padding_mask = create_padding_mask(inp)
tar_padding_mask = create_padding_mask(tar)
look_ahead_mask = create_look_ahead_mask(tar.shape[1])
combined_mask = tf.math.maximum(tar_padding_mask, look_ahead_mask)

# 初始化一个 Decoder
decoder = Decoder(num_layers, d_model, num_heads, dff, target_vocab_size)

# 将 2 维的索引序列以及遮罩丢入 Decoder
print("tar:", tar)
print("-" * 20)
print("combined_mask:", combined_mask)
print("-" * 20)
print("enc_out:", enc_out)
print("-" * 20)
print("inp_padding_mask:", inp_padding_mask)
print("-" * 20)
dec_out, attn = decoder(tar, enc_out, training=False, 
                        combined_mask=combined_mask,
                        inp_padding_mask=inp_padding_mask)
print("dec_out:", dec_out)
print("-" * 20)
for block_name, attn_weights in attn.items():
    print(f"{block_name}.shape: {attn_weights.shape}")

tar: tf.Tensor(
[[4205   10  241   86   27    3 4206    0    0    0]
 [4205  165  489  398  191   14    7  560    3 4206]], shape=(2, 10), dtype=int64)
--------------------
combined_mask: tf.Tensor(
[[[[0. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
   [0. 0. 1. 1. 1. 1. 1. 1. 1. 1.]
   [0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
   [0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
   [0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 1. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
   [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]]]

 [[[0. 1. 1. 1. 1. 1. 1. 1. 1. 1.]
    [0. 0. 1. 1. 1. 1. 1. 1. 1. 1.]
    [0. 0. 0. 1. 1. 1. 1. 1. 1. 1.]
    [0. 0. 0. 0. 1. 1. 1. 1. 1. 1.]
    [0. 0. 0. 0. 0. 1. 1. 1. 1. 1.]
    [0. 0. 0. 0. 0. 0. 1. 1. 1. 1.]
    [0. 0. 0. 0. 0. 0. 0. 1. 1. 1.]
    [0. 0. 0. 0. 0. 0. 0. 0. 1. 1.]
    [0. 0. 0. 0. 0. 0. 0. 0. 0. 1.]
    [0. 0. 0. 0. 0. 0. 0. 0. 0. 0.]]]], shape=(2, 1, 10, 10), dtype=float32)
   --------------------
enc_out: tf.Tensor(
[[[-0.7849332  -0.5919684  -0.33270505  1.7096066 ]
  [-0.5070654  -0.5110136  -0.7082318   1.726311  ]
  [-0.39270183 -0.03102639 -1.158362    1.5820901 ]
  [-0.5561629   0.38050282 -1.2407898   1.4164499 
  [-0.90432     0.19381052 -0.8472892   1.5577985 ]
  [-0.97321564 -0.22992788 -0.4652462   1.6683896 ]
  [-0.84681976 -0.5434473  -0.31013608  1.7004032 ]
  [-0.62432766 -0.56790507 -0.539001    1.7312336 ]]

 [[-0.77423775 -0.6076471  -0.32800597  1.7098908 ]
  [-0.47978252 -0.5615605  -0.68602914  1.7273722 ]
  [-0.30068305 -0.07366991 -1.1973959   1.5717487 ]
  [-0.5147841   0.2787246  -1.2290851   1.4651446 ]
  [-0.89634496  0.2675462  -0.8954112   1.52421   ]
  [-0.97553635 -0.22618684 -0.4656965   1.6674198 ]
  [-0.87600434 -0.5448401  -0.27099532  1.6918398 ]
  [-0.60130465 -0.5993665  -0.5306774   1.7313484 ]]], shape=(2, 8, 4), dtype=float32)
--------------------
inp_padding_mask: tf.Tensor(
[[[[0. 0. 0. 0. 0. 0. 1. 1.]]]
 [[[0. 0. 0. 0. 0. 0. 0. 0.]]]], shape=(2, 1, 1, 8), dtype=float32)
--------------------
dec_out: tf.Tensor(
[[[-0.5652141  -1.0581813   1.600075    0.02332011]
  [-0.34019774 -1.2377603   1.5330346   0.04492359]
  [ 0.3675252  -1.4228352   1.3287866  -0.2734765 ]
  [ 0.09472068 -1.353683    1.4559422  -0.19697984]
  [-0.38392055 -1.0940721   1.6231282  -0.14513558]
  [-0.41729763 -1.0276326   1.6514215  -0.20649135]
  [-0.3302343  -1.0454822   1.6500466  -0.27433014]
  [-0.1923209  -1.1254803   1.6149355  -0.29713422]
  [ 0.40822834 -1.3586452   1.3515034  -0.40108633]
  [ 0.19979587 -1.4183372   1.3857942  -0.1672527 ]]

 [[-0.56504554 -1.054449    1.602678    0.01681651]
  [-0.36043385 -1.2348608   1.5300139   0.0652808 ]
  [ 0.24521776 -1.4295446   1.3651297  -0.18080302]
  [-0.06483467 -1.3449186   1.4773033  -0.06755   ]
  [-0.41885287 -1.0775515   1.6267892  -0.1303851 ]
  [-0.40018192 -1.0338533   1.6504982  -0.21646297]
  [-0.3531929  -1.0375831   1.6523482  -0.26157203]
  [-0.24463172 -1.1371143   1.6107951  -0.22904922]
  [ 0.19615419 -1.362728    1.4271017  -0.2605278 ]
  [ 0.08419974 -1.3687493   1.4467624  -0.16221291]]], shape=(2, 10, 4), dtype=float32)
--------------------
decoder_layer1_block1.shape: (2, 2, 10, 10)
decoder_layer1_block2.shape: (2, 2, 10, 8)
decoder_layer2_block1.shape: (2, 2, 10, 10)
decoder_layer2_block2.shape: (2, 2, 10, 8)

麻雀虽小，五脏俱全。虽然我们是使用 demo 数据，但基本上这就是你在呼叫 Decoder 时需要做的所有事情：

初始时给它中文（目标语言）的字典大小、其他超参数
输入中文 batch 的索引序列
也要输入两个遮罩以及 Encoder 输出enc_out

Decoder 的输出你现在应该都可以很轻松地解读才是。基本上跟 Decoder layer 一模一样，只差在我们额外输出一个 Python dict，里头存放所有 Decoder layers 的注意权重。

6.7 第一个 Transformer

没错，终于到了这个时刻。在实现 Transformer 之前先点击影片来简单回顾一下我们在这一章实现了什么些玩意儿：

transformer-imple.gif

Transformer 本身只有 3 个 layers

在我们前面已经将大大小小的 layers 一一实作并组装起来以后，真正的 Transformer 模型只需要 3 个元件：

Encoder
Decoder
Final linear layer

马上让我们看看 Transformer 的实现：

# Transformer 之上已經沒有其他 layers 了，我們使用 tf.keras.Model 建立一個模型
class Transformer(tf.keras.Model):
    # 初始參數包含 Encoder & Decoder 都需要超參數以及中英字典數目
    def __init__(self, num_layers, d_model, num_heads, dff, input_vocab_size, target_vocab_size, rate=0.1):
        super(Transformer, self).__init__()

        self.encoder = Encoder(num_layers, d_model, num_heads, dff, input_vocab_size, rate)

        self.decoder = Decoder(num_layers, d_model, num_heads, dff, target_vocab_size, rate)
        # 這個 FFN 輸出跟中文字典一樣大的 logits 數，等通過 softmax 就代表每個中文字的出現機率
        self.final_layer = tf.keras.layers.Dense(target_vocab_size)
  
    # enc_padding_mask 跟 dec_padding_mask 都是英文序列的 padding mask，
    # 只是一個給 Encoder layer 的 MHA 用，一個是給 Decoder layer 的 MHA 2 使用
    def call(self, inp, tar, training, enc_padding_mask, combined_mask, dec_padding_mask):
        enc_output = self.encoder(inp, training, enc_padding_mask)  # (batch_size, inp_seq_len, d_model)
    
        # dec_output.shape == (batch_size, tar_seq_len, d_model)
        dec_output, attention_weights = self.decoder(tar, enc_output, training, combined_mask, dec_padding_mask)
    
        # 將 Decoder 輸出通過最後一個 linear layer
        final_output = self.final_layer(dec_output)  # (batch_size, tar_seq_len, target_vocab_size)
    
        return final_output, attention_weights

扣掉注解，Transformer 的实现本身非常简短。

被输入Transformer 的多个2 维英文张量inp 会一路通过Encoder 里头的词嵌入层，位置编码以及N 个Encoder layers 后被转换成Encoder 输出enc_output，接着对应的中文序列tar 则会在Decoder 里头走过相似的旅程并在每一层的Decoder layer 利用MHA 2 关注Encoder 的输出enc_output，最后被Decoder 输出。

而 Decoder 的输出 dec_output 则会通过 Final linear layer，被转成进入 Softmax 前的 logits final_output，其 logit 的数目则跟中文字典里的子词数相同。

因为Transformer 把Decoder 也包起来了，现在我们连Encoder 输出enc_output也不用管，只要把英文（来源）以及中文（目标）的索引序列batch 丢入Transformer，它就会输出最后一维为中文字典大小的张量。第 2 维是输出序列，里头每一个位置的向量就代表着该位置的中文字的概率分布（事实上通过 softmax 才是，但这边先这样说方便你理解）：

输入：
- 英文序列：（batch_size, inp_seq_len）
- 中文序列：（batch_size, tar_seq_len）
输出：
- 生成序列：（batch_size, tar_seq_len, target_vocab_size）
- 注意权重的 dict

让我们马上建一个 Transformer，并假设我们已经准备好用 demo 数据来训练它做英翻中：

# 超参数
num_layers = 1
d_model = 4
num_heads = 2
dff = 8

# + 2 是为了 <start> & <end> token
input_vocab_size = subword_encoder_en.vocab_size + 2
output_vocab_size = subword_encoder_zh.vocab_size + 2

# 重点中的重点。训练时用前一个字来预测下一个中文字
tar_inp = tar[:, :-1]
tar_real = tar[:, 1:]

# 来源 / 目标语言用的遮罩。注意 `comined_mask` 已经将目标语言的两种遮罩合而为一
inp_padding_mask = create_padding_mask(inp)
tar_padding_mask = create_padding_mask(tar_inp)
look_ahead_mask = create_look_ahead_mask(tar_inp.shape[1])
combined_mask = tf.math.maximum(tar_padding_mask, look_ahead_mask)

# 初始化我们的第一个 transformer
transformer = Transformer(num_layers, d_model, num_heads, dff, 
                          input_vocab_size, output_vocab_size)

# 将英文、中文序列丢入取得 Transformer 预测下个中文字的结果
predictions, attn_weights = transformer(inp, tar_inp, False, inp_padding_mask, 
                                        combined_mask, inp_padding_mask)

print("tar:", tar)
print("-" * 20)
print("tar_inp:", tar_inp)
print("-" * 20)
print("tar_real:", tar_real)
print("-" * 20)
print("predictions:", predictions)

tar: tf.Tensor(
[[4205   10  241   86   27    3 4206    0    0    0]
 [4205  165  489  398  191   14    7  560    3 4206]], shape=(2, 10), dtype=int64)
--------------------
tar_inp: tf.Tensor(
[[4205   10  241   86   27    3 4206    0    0]
 [4205  165  489  398  191   14    7  560    3]], shape=(2, 9), dtype=int64)
--------------------
tar_real: tf.Tensor(
[[  10  241   86   27    3 4206    0    0    0]
 [ 165  489  398  191   14    7  560    3 4206]], shape=(2, 9), dtype=int64)
--------------------
predictions: tf.Tensor(
[[[ 0.01349578 -0.00199539 -0.00217387 ... -0.03862738 -0.03212879
   -0.07692747]
  [ 0.037483    0.01585471 -0.02548708 ... -0.04276202 -0.02495992
   -0.05491882]
  [ 0.05718528  0.0288353  -0.04577483 ... -0.0450176  -0.01315334
   -0.03639907]
  ...
  [ 0.01202047 -0.00400385 -0.00099438 ... -0.03859971 -0.03085513
   -0.0797975 ]
  [ 0.0235797   0.00501019 -0.01193091 ... -0.04091505 -0.02892826
   -0.06939011]
  [ 0.04867784  0.02382022 -0.03683803 ... -0.04392421 -0.01941058
   -0.04347047]]

 [[ 0.01676657 -0.00080312 -0.00556347 ... -0.03981712 -0.02937311
   -0.07665333]
  [ 0.03873826  0.01607161 -0.02685272 ... -0.04328423 -0.02345929
   -0.05522631]
  [ 0.0564083   0.02865588 -0.04492006 ... -0.04475704 -0.014088
   -0.03639095]
  ...
  [ 0.01514172 -0.00298804 -0.00426158 ... -0.03976889 -0.02800199
   -0.07974622]
  [ 0.02867933  0.00800282 -0.01704068 ... -0.04215823 -0.02618418
   -0.06638923]
  [ 0.05056309  0.02489874 -0.03880978 ... -0.04421616 -0.01803544
   -0.04204436]]], shape=(2, 9, 4207), dtype=float32)

有了前面的各种 layers，建立一个 Transformer 并不难。但要输入什么数据就是一门大学问了：

...
tar_inp = tar[:, :-1]
tar_real = tar[:, 1:]
predictions, attn_weights = transformer(inp, tar_inp, False, ...)
...

为何是丢少了尾巴一个字的 tar_inp 序列进去 Transformer，而不是直接丢 tar 呢？

别忘记我们才刚初始一个 Transformer，里头所有 layers 的权重都是随机的，你可不能指望它真的会什么「黑魔法」来帮你翻译。我们得先训练才行。但训练时如果你把整个正确的中文序列 tar都进去给 Transformer 看，你期待它产生什么？一首新的中文诗吗？

如果你曾经实现过序列生成模型或是看过我之前的语言模型文章，就会知道在序列生成任务里头，模型获得的正确答案是输入序列往左位移一个位置的结果。

这样讲很抽象，让我们看个影片了解序列生成是怎么运作的：

了解序列生成以及如何训练一个生成模型

你现在应该明白 Transformer 在训练的时候并不是吃进去整个中文序列，而是吃进去一个去掉尾巴的序列 tar_inp，然后试着去预测「左移」一个字以后的序列 tar_real。同样概念当然也可以运用到以 RNN 或是 CNN-based 的模型上面。

从影片中你也可以发现给定 tar_inp 序列中的任一位置，其对应位置的 tar_real 就是下个时间点模型应该要预测的中文字。

序列生成任务可以被视为是一个分类任务（Classification），而每一个中文字都是一个分类。而 Transformer 就是要去产生一个中文字的概率分布，想办法跟正解越接近越好。

跟用已训练的Transformer 做预测时不同，在训练时为了稳定模型表现，我们并不会将Transformer 的输出再度丢回去当做其输入（人形蜈蚣？），而是像影片中所示，给它左移一个位置后的序列tar_real 当作正解让它去最小化error。

这种无视模型预测结果，而将正确解答丢入的训练方法一般被称作 teacher forcing。你也可以参考教授的 Sequence-to-sequence Learning 教学。

7. 定义损失函数与指标

因为被视为是一个分类任务，我们可以使用 cross entropy 来计算序列生成任务中实际的中文字跟模型预测的中文字分布（distribution）相差有多远。

这边简单定义一个损失函数：

loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')

# 假设我们要解的是一个 binary classifcation， 0 跟 1 个代表一个 label
real = tf.constant([1, 1, 0], shape=(1, 3), dtype=tf.float32)
pred = tf.constant([[0, 1], [0, 1], [0, 1]], dtype=tf.float32)
loss_object(real, pred)

<tf.Tensor: shape=(3,), dtype=float32, numpy=array([0.31326166, 0.31326166, 1.3132616 ], dtype=float32)>

如果你曾做过分类问题，应该能看出预测序列pred 里头的第 3 个预测结果出错因此 entropy 值上升。损失函数loss_object 做的事情就是比较 2 个序列并计算 cross entropy：

real：一个包含 N 个正确 labels 的序列
pred：一个包含 N 个维度为 label 数的 logit 序列

我们在这边将 reduction 参数设为 none，请loss_object 不要把每个位置的 error 加总。而这是因为我们之后要自己把 <pad> token 出现的位置的损失舍弃不计。

而将 from_logits 参数设为True 是因为从 Transformer 得到的预测还没有经过 softmax，因此加和还不等于 1：

print("predictions:", predictions)
print("-" * 20)
print(tf.reduce_sum(predictions, axis=-1))

predictions: tf.Tensor(
[[[ 0.01349578 -0.00199539 -0.00217387 ... -0.03862738 -0.03212879
   -0.07692747]
  [ 0.037483    0.01585471 -0.02548708 ... -0.04276202 -0.02495992
   -0.05491882]
  [ 0.05718528  0.0288353  -0.04577483 ... -0.0450176  -0.01315334
   -0.03639907]
  ...
  [ 0.01202047 -0.00400385 -0.00099438 ... -0.03859971 -0.03085513
   -0.0797975 ]
  [ 0.0235797   0.00501019 -0.01193091 ... -0.04091505 -0.02892826
   -0.06939011]
  [ 0.04867784  0.02382022 -0.03683803 ... -0.04392421 -0.01941058
   -0.04347047]]

 [[ 0.01676657 -0.00080312 -0.00556347 ... -0.03981712 -0.02937311
   -0.07665333]
  [ 0.03873826  0.01607161 -0.02685272 ... -0.04328423 -0.02345929
   -0.05522631]
  [ 0.0564083   0.02865588 -0.04492006 ... -0.04475704 -0.014088
   -0.03639095]
  ...
  [ 0.01514172 -0.00298804 -0.00426158 ... -0.03976889 -0.02800199
   -0.07974622]
  [ 0.02867933  0.00800282 -0.01704068 ... -0.04215823 -0.02618418
   -0.06638923]
  [ 0.05056309  0.02489874 -0.03880978 ... -0.04421616 -0.01803544
   -0.04204436]]], shape=(2, 9, 4207), dtype=float32)
--------------------
tf.Tensor(
[[1.3761909 2.9352095 3.8687317 3.4191105 2.608357  1.5664345 1.1489892
  1.9882674 3.5525477]
 [1.4309797 2.9219136 3.873899  3.5009165 2.6499162 1.6611676 1.1839213
  2.2150593 3.6206641]], shape=(2, 9), dtype=float32)

有了 loss_object 实际算 cross entropy 以后，我们需要另外一个函数来建立遮罩并加总序列里头不包含 token位置的损失：

def loss_function(real, pred):
  # 这次的 mask 将序列中不等于 0 的位置视为 1，其余为 0
  mask = tf.math.logical_not(tf.math.equal(real, 0))
  # 照样计算所有位置的 cross entropy 但不加总
  loss_ = loss_object(real, pred)
  mask = tf.cast(mask, dtype=loss_.dtype)
  loss_ *= mask  # 只计算非 <pad> 位置的损失
  
  return tf.reduce_mean(loss_)

我另外再定义两个 tf.keras.metrics，方便之后使用 TensorBoard 来追踪模型 performance：

train_loss = tf.keras.metrics.Mean(name='train_loss')
train_accuracy = tf.keras.metrics.SparseCategoricalAccuracy(
    name='train_accuracy')

8. 设置超参数

前面实现了那么多 layers，你应该还记得有哪些是你自己可以调整的超参数吧？

让我帮你全部列出来：

num_layers 决定 Transfomer 里头要有几个 Encoder / Decoder layers
d_model 决定我们子词的 representation space 维度
num_heads 要做几头的自注意力运算
dff 决定 FFN 的中间维度
dropout_rate 预设 0.1，一般用预设值即可
input_vocab_size：输入语言（英文）的字典大小
target_vocab_size：输出语言（中文）的字典大小

论文里头最基本的 Transformer 配置为：

num_layers=6
d_model=512
dff=2048

有大量数据以及大的 Transformer，你可以在很多机器学习任务都达到不错的成绩。为了不要让训练时间太长，在这篇文章里头我会把 Transformer 里头的超参数设小一点：

num_layers = 4 
d_model = 128
dff = 512
num_heads = 8

input_vocab_size = subword_encoder_en.vocab_size + 2
target_vocab_size = subword_encoder_zh.vocab_size + 2
dropout_rate = 0.1  # 预设值

print("input_vocab_size:", input_vocab_size)
print("target_vocab_size:", target_vocab_size)

input_vocab_size: 8115
target_vocab_size: 4207

4 层 Encoder / Decoder layers 不算贪心，小巫见大巫（笑

9. 设置 Optimizer

我们在这边跟论文一致，使用 Adam optimizer 以及自定义的 learning rate scheduler：

lr-equation.jpg

这 schedule 让训练过程的前 warmup_steps 的 learning rate 线性增加，在那之后则跟步骤数 step_num的反平方根成比例下降。不用担心你没有完全理解这公式，我们一样可以直接使用 TensorFlow 官方教学的实现：

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    # 论文预设 `warmup_steps` = 4000
    def __init__(self, d_model, warmup_steps=4000):
        super(CustomSchedule, self).__init__()
    
        self.d_model = d_model
        self.d_model = tf.cast(self.d_model, tf.float32)

        self.warmup_steps = warmup_steps

    def __call__(self, step):
        arg1 = tf.math.rsqrt(step)
        arg2 = step * (self.warmup_steps ** -1.5)
        return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)
    
# 将客制化 learning rate schdeule 丢入 Adam opt.
# Adam opt. 的参数都跟论文相同
learning_rate = CustomSchedule(d_model)
optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, 
                                     epsilon=1e-9)

我们可以观察看看这个 schedule 是怎么随着训练步骤而改变 learning rate 的：

d_models = [128, 256, 512]
warmup_steps = [1000 * i for i in range(1, 4)]

schedules = []
labels = []
colors = ["blue", "red", "black"]
for d in d_models:
  schedules += [CustomSchedule(d, s) for s in warmup_steps]
  labels += [f"d_model: {d}, warm: {s}" for s in warmup_steps]

for i, (schedule, label) in enumerate(zip(schedules, labels)):
  plt.plot(schedule(tf.range(10000, dtype=tf.float32)), 
           label=label, color=colors[i // 3])

plt.legend()

plt.ylabel("Learning Rate")
plt.xlabel("Train Step")

Text(0.5, 0, 'Train Step')

不同 d_model 以及 warmup_steps 的 learning rate 变化

你可以明显地看到所有 schedules 都先经过 warmup_steps 个步骤直线提升 learning rate，接着逐渐平滑下降。另外我们也会给比较高维的 d_model 维度比较小的 learning rate。

10. 实际训练以及定时存档

好啦，什么都准备齐全了，让我们开始训练 Transformer 吧！记得使用前面已经定义好的超参数来初始化一个全新的 Transformer：

transformer = Transformer(num_layers, d_model, num_heads, dff,
                          input_vocab_size, target_vocab_size, dropout_rate)

print(f"""这个 Transformer 有 {num_layers} 层 Encoder / Decoder layers
d_model: {d_model}
num_heads: {num_heads}
dff: {dff}
input_vocab_size: {input_vocab_size}
target_vocab_size: {target_vocab_size}
dropout_rate: {dropout_rate}

""")

这个 Transformer 有 4 层 Encoder / Decoder layers
d_model: 128
num_heads: 8
dff: 512
input_vocab_size: 8115
target_vocab_size: 4207
dropout_rate: 0.1

打游戏时你会记得要定期存档以防任何意外发生，训练深度学习模型也是同样道理。设置 checkpoint 来定期储存 / 读取模型及 optimizer 是必备的。

我们在底下会定义一个 checkpoint 路径，此路径包含了各种超参数的信息，方便之后比较不同实验的结果并载入已训练的进度。我们也需要一个 checkpoint manager 来做所有跟存读模型有关的杂事，并只保留最新 5 个 checkpoints 以避免占用太多空间：

# 方便比较不同实验/ 不同超参数设定的结果
run_id = f"{num_layers}layers_{d_model}d_{num_heads}heads_{dff}dff_{train_perc}train_perc"
checkpoint_path = os.path.join(checkpoint_path, run_id)
log_dir = os.path.join(log_dir, run_id)

# tf.train.Checkpoint 可以帮我们把想要存下来的东西整合起来，方便储存与读取
# 一般来说你会想存下模型以及 optimizer 的状态
ckpt = tf.train.Checkpoint(transformer=transformer,
                           optimizer=optimizer)

# ckpt_manager 会去 checkpoint_path 看有没有符合 ckpt 里头定义的东西
# 存档的时候只保留最近 5 次 checkpoints，其他自动删除
ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

# 如果在 checkpoint 路径上有发现档案就读进来
if ckpt_manager.latest_checkpoint:
    ckpt.restore(ckpt_manager.latest_checkpoint)
  
    # 如果在 checkpoint 路径上有发现档案就读进来
    last_epoch = int(ckpt_manager.latest_checkpoint.split("-")[-1])
    print(f'已读取最新的 checkpoint，模型已训练 {last_epoch} epochs。')
else:
    last_epoch = 0
    print("沒找到 checkpoint，从头训练。")

沒找到 checkpoint，从头训练。

我知道你在想什么。

「诶！？你不当场训练吗？」「直接载入已训练的模型太狗了吧！」

拜托，我都训练 N 遍了，每次都重新训练也太没意义了。而且你能想像为了写一个章节我就得重新训练一个 Transformer 来 demo 吗？这样太没效率了。比起每次重新训练模型，这才是你在真实世界中应该做的事情：尽可能恢复之前的训练进度来节省时间。

不过放心，我仍会秀出完整的训练代码让你可以执行第一次的训练。当你想要依照本文训练自己的 Transformer 时会感谢有 checkpoint manager 的存在。现在假设我们还没有 checkpoints。

在实际训练 Transformer 之前还需要定义一个简单函数来产生所有的遮罩：

# 为 Transformer 的 Encoder / Decoder 准备遮罩
def create_masks(inp, tar):
  # 英文句子的 padding mask，要交給 Encoder layer 自注意力機制用的
  enc_padding_mask = create_padding_mask(inp)
  
  # 同樣也是英文句子的 padding mask，但是是要交給 Decoder layer 的 MHA 2 
  # 關注 Encoder 輸出序列用的
  dec_padding_mask = create_padding_mask(inp)
  
  # Decoder layer 的 MHA1 在做自注意力機制用的
  # `combined_mask` 是中文句子的 padding mask 跟 look ahead mask 的疊加
  look_ahead_mask = create_look_ahead_mask(tf.shape(tar)[1])
  dec_target_padding_mask = create_padding_mask(tar)
  combined_mask = tf.maximum(dec_target_padding_mask, look_ahead_mask)
  
  return enc_padding_mask, combined_mask, dec_padding_mask

如果没有本文前面针对遮罩的详细说明，很多第一次实现的人得花不少时间来确实地掌握这些遮罩的用途。不过对现在的你来说应该也是小菜一碟。

一个数据集包含多个 batch，而每次拿一个 batch 来训练的步骤就称作 train_step。为了让程式码更简洁以及容易优化，我们会定义 Transformer 在一次训练步骤（处理一个 batch）所需要做的所有事情。

不限于 Transformer，一般来说 train_step 函数里会有几个重要步骤：

对训练数据做些必要的前处理
将数据丢入模型，取得预测结果
用预测结果跟正确解答计算 loss
取出梯度并利用 optimizer 做梯度下降

有了这个概念以后看看代码：

@tf.function  # 让 TensorFlow 帮我们将 eager code 优化并加快运算
def train_step(inp, tar):
    # 前面说过的，用去尾的原始序列去预测下一个字的序列
    tar_inp = tar[:, :-1]
    tar_real = tar[:, 1:]
  
    # 建立 3 个遮罩
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)
  
    # 纪录 Transformer 的所有运算过程以方便之后做梯度下降
    with tf.GradientTape() as tape:
        # 注意是丢入 `tar_inp` 而非 `tar`。记得将 `training` 参数设定为 True
        predictions, _ = transformer(inp, tar_inp, 
                                     True, 
                                     enc_padding_mask, 
                                     combined_mask, 
                                     dec_padding_mask)
        # 跟影片中显示的相同，计算左移一个字的序列跟模型预测分布之间的差异，当作 loss
        loss = loss_function(tar_real, predictions)

    # 取出梯度并呼叫前面定义的 Adam optimizer 帮我们更新 Transformer 里头可训练的参数
    gradients = tape.gradient(loss, transformer.trainable_variables)    
    optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))
  
    # 将 loss 以及训练 acc 记录到 TensorBoard 上，非必要
    train_loss(loss)
    train_accuracy(tar_real, predictions)

如果你曾经以TensorFlow 2 实现过稍微复杂一点的模型，应该就知道 train_step函数的写法非常固定：

对输入数据做些前处理（本文中的遮罩、将输出序列左移当成正解 etc.）
利用 tf.GradientTape 轻松记录数据被模型做的所有转换并计算 loss
将梯度取出并让 optimzier 对可被训练的权重做梯度下降（上升）

你完全可以用一模一样的方式将任何复杂模型的处理过程包在train_step 函数，这样可以让我们之后在 iterate 数据集时非常轻松。而且最重要的是可以用 tf.function 来提高此函数里头运算的速度。你可以点击连结来了解更多。

处理一个 batch 的 train_step 函数也有了，就只差写个 for loop 将数据集跑个几遍了。我之前的模型虽然训练了 50 个 epochs，但事实上大概 30 epochs 翻译的结果就差不多稳定了。所以让我们将 EPOCHS 设定为 30：

# 定义我们要看几遍数据集
EPOCHS = 30
print(f"此超参数组合的 Transformer 已经训练 {last_epoch} epochs。")
print(f"剩余 epochs：{min(0, last_epoch - EPOCHS)}")


# 用来写资讯到 TensorBoard，非必要但十分推荐
summary_writer = tf.summary.create_file_writer(log_dir)

# 比对设定的 `EPOCHS` 以及已训练的 `last_epoch` 来决定还要训练多少 epochs
for epoch in range(last_epoch, EPOCHS):
    start = time.time()
  
    # 重置纪录 TensorBoard 的 metrics
    train_loss.reset_states()
    train_accuracy.reset_states()
  
    # 一个 epoch 就是把我们定义的训练资料集一个一个 batch 拿出来处理，直到看完整个数据集
    for (step_idx, (inp, tar)) in enumerate(train_dataset):
        # 每次 step 就是将数据丢入 Transformer，让它生预测结果并计算梯度最小化 loss
        train_step(inp, tar)  

    # 每个 epoch 完成就存一次档
    if (epoch + 1) % 1 == 0:
        ckpt_save_path = ckpt_manager.save()
        print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,ckpt_save_path))
    
    # 将 loss 以及 accuracy 写到 TensorBoard 上
    with summary_writer.as_default():
        tf.summary.scalar("train_loss", train_loss.result(), step=epoch + 1)
        tf.summary.scalar("train_acc", train_accuracy.result(), step=epoch + 1)
  
    print('Epoch {} Loss {:.4f} Accuracy {:.4f}'.format(epoch + 1, 
                                                        train_loss.result(), 
                                                        train_accuracy.result()))
    print('Time taken for 1 epoch: {} secs\n'.format(time.time() - start))

此超参数组合的 Transformer 已经训练 0 epochs。
剩余 epochs：-30
Saving checkpoint for epoch 1 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-1
Epoch 1 Loss 5.1843 Accuracy 0.0219
Time taken for 1 epoch: 89.55020833015442 secs

Saving checkpoint for epoch 2 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-2
Epoch 2 Loss 4.2425 Accuracy 0.0604
Time taken for 1 epoch: 21.873889207839966 secs

Saving checkpoint for epoch 3 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-3
Epoch 3 Loss 3.7423 Accuracy 0.0987
Time taken for 1 epoch: 21.901566743850708 secs

Saving checkpoint for epoch 4 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-4
Epoch 4 Loss 3.2644 Accuracy 0.1512
Time taken for 1 epoch: 22.083024501800537 secs

Saving checkpoint for epoch 5 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-5
Epoch 5 Loss 2.9634 Accuracy 0.1810
Time taken for 1 epoch: 22.050684452056885 secs

Saving checkpoint for epoch 6 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-6
Epoch 6 Loss 2.7756 Accuracy 0.1988
Time taken for 1 epoch: 25.719687461853027 secs

Saving checkpoint for epoch 7 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-7
Epoch 7 Loss 2.6346 Accuracy 0.2122
Time taken for 1 epoch: 22.85287618637085 secs

Saving checkpoint for epoch 8 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-8
Epoch 8 Loss 2.5183 Accuracy 0.2242
Time taken for 1 epoch: 18.721409797668457 secs

Saving checkpoint for epoch 9 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-9
Epoch 9 Loss 2.4171 Accuracy 0.2353
Time taken for 1 epoch: 18.663178205490112 secs

Saving checkpoint for epoch 10 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-10
Epoch 10 Loss 2.3204 Accuracy 0.2458
Time taken for 1 epoch: 25.891611576080322 secs

Saving checkpoint for epoch 11 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-11
Epoch 11 Loss 2.2223 Accuracy 0.2573
Time taken for 1 epoch: 18.789816856384277 secs

Saving checkpoint for epoch 12 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-12
Epoch 12 Loss 2.1319 Accuracy 0.2685
Time taken for 1 epoch: 22.33806586265564 secs

Saving checkpoint for epoch 13 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-13
Epoch 13 Loss 2.0458 Accuracy 0.2796
Time taken for 1 epoch: 18.877813816070557 secs

Saving checkpoint for epoch 14 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-14
Epoch 14 Loss 1.9643 Accuracy 0.2912
Time taken for 1 epoch: 18.858903884887695 secs

Saving checkpoint for epoch 15 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-15
Epoch 15 Loss 1.8875 Accuracy 0.3020
Time taken for 1 epoch: 18.890562295913696 secs

Saving checkpoint for epoch 16 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-16
Epoch 16 Loss 1.8178 Accuracy 0.3120
Time taken for 1 epoch: 22.47147297859192 secs

Saving checkpoint for epoch 17 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-17
Epoch 17 Loss 1.7531 Accuracy 0.3211
Time taken for 1 epoch: 18.98854422569275 secs

Saving checkpoint for epoch 18 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-18
Epoch 18 Loss 1.6899 Accuracy 0.3305
Time taken for 1 epoch: 18.987966775894165 secs

Saving checkpoint for epoch 19 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-19
Epoch 19 Loss 1.6200 Accuracy 0.3406
Time taken for 1 epoch: 18.95727038383484 secs

Saving checkpoint for epoch 20 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-20
Epoch 20 Loss 1.5555 Accuracy 0.3499
Time taken for 1 epoch: 18.99857258796692 secs

Saving checkpoint for epoch 21 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-21
Epoch 21 Loss 1.4968 Accuracy 0.3590
Time taken for 1 epoch: 19.01795792579651 secs

Saving checkpoint for epoch 22 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-22
Epoch 22 Loss 1.4447 Accuracy 0.3668
Time taken for 1 epoch: 19.078711986541748 secs

Saving checkpoint for epoch 23 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-23
Epoch 23 Loss 1.3984 Accuracy 0.3738
Time taken for 1 epoch: 19.144370317459106 secs

Saving checkpoint for epoch 24 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-24
Epoch 24 Loss 1.3535 Accuracy 0.3805
Time taken for 1 epoch: 19.05727791786194 secs

Saving checkpoint for epoch 25 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-25
Epoch 25 Loss 1.3142 Accuracy 0.3866
Time taken for 1 epoch: 22.631419897079468 secs

Saving checkpoint for epoch 26 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-26
Epoch 26 Loss 1.2765 Accuracy 0.3926
Time taken for 1 epoch: 19.017268657684326 secs

Saving checkpoint for epoch 27 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-27
Epoch 27 Loss 1.2441 Accuracy 0.3969
Time taken for 1 epoch: 19.065359115600586 secs

Saving checkpoint for epoch 28 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-28
Epoch 28 Loss 1.2106 Accuracy 0.4023
Time taken for 1 epoch: 19.06916570663452 secs

Saving checkpoint for epoch 29 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-29
Epoch 29 Loss 1.1835 Accuracy 0.4068
Time taken for 1 epoch: 19.07039451599121 secs

Saving checkpoint for epoch 30 at nmt/checkpoints/4layers_128d_8heads_512dff_20train_perc/ckpt-30
Epoch 30 Loss 1.1560 Accuracy 0.4107
Time taken for 1 epoch: 19.10555648803711 secs

如信息所示，当指定的 EPOCHS「落后」于之前的训练进度我们就不再训练了。但如果是第一次训练或是训练到指定 EPOCHS的一部分，我们都会从正确的地方开始训练并存档，不会浪费到训练时间或计算资源。

这边的逻辑也很简单，在每个 epoch 都：

（非必要）重置写到 TensorBoard 的 metrics 的值
将整个数据集的 batch 取出，交给 train_step 函数处理
（非必要）存 checkpoints
（非必要）将当前 epoch 结果写到 TensorBoard
（非必要）在标准输出显示当前 epoch 结果

是的，如果你真的只是想要训练个模型，什么其他事情都不想考虑的话那你可以：

# 87 分，不能再高了。
for epoch in range(EPOCHS):
  for inp, tar in train_dataset:
    train_step(inp, tar)

嗯 ... 话是这么说，但我仍然建议你至少要记得存档并将训练过程显示出来。

编者按：我是在Google的Colab Notebooks中进行的训练，在这个计算能力下，我们定义的 4 层 Transformer 大约每 19 秒就可以看完一遍有 3 万笔训练例子的数据集，而且你从上面的 loss 以及 accuracy 可以看出来 Transformer 至少在训练集里头进步地挺快的。

而就我自己的观察大约经过 30 个 epochs 翻译结果就很稳定了。所以你大约只需半个小时就能有一个非常简单，有点水准的英翻中 Transformer（在至少有个一般 GPU 的情况）。

但跟看上面的 log 比起来，我个人还是比较推荐使用 TensorBoard。在 TensorFlow 2 里头，你甚至能直接在 Jupyter Notebook 或是 Colab 里头开启它：

%load_ext tensorboard
%tensorboard --logdir {log_dir}

<IPython.core.display.Javascript object>

使用 TensorBoard 可以让你轻松比较不同超参数的训练结果

透过 TensorBoard，你能非常清楚地比较不同实验以及不同点子的效果，知道什么 work 什么不 work，进而修正之后尝试的方向。如果只是简单写个print，那你永远只会看到最新一次训练过程的 log，然后忘记之前到底发生过什么事。

11. 实际进行英翻中

有了已经训练一阵子的 Transformer，当然得拿它来实际做做翻译。

跟训练的时候不同，在做预测时我们不需做 teacher forcing 来稳定 Transformer 的训练过程。反之，我们将 Transformer 在每个时间点生成的中文索引加到之前已经生成的序列尾巴，并以此新序列作为其下一次的输入。这是因为 Transformer 事实上是一个自回归模型（Auto-regressive model）：依据自己生成的结果预测下次输出。

利用 Transformer 进行翻译（预测）的逻辑如下：

将输入的英文句子利用 Subword Tokenizer 转换成子词索引序列（还记得 inp 吧？）
在该英文索引序列前后加上代表英文 BOS / EOS 的tokens
在 Transformer 输出序列长度达到 MAX_LENGTH 之前重复以下步骤：
- 为目前已经生成的中文索引序列产生新的遮罩
- 将刚刚的英文序列、当前的中文序列以及各种遮罩放入 Transformer
- 将 Transformer 输出序列的最后一个位置的向量取出，并取 argmax 取得新的预测中文索引
- 将此索引加到目前的中文索引序列里头作为 Transformer 到此为止的输出结果
- 如果新生成的中文索引为 <end> 则代表中文翻译已全部生成完毕，直接回传
将最后得到的中文索引序列回传作为翻译结果

是的，一个时间点生成一个中文字，而在第一个时间点因为 Transformer 还没有任何输出，我们会丢中文字的 <start> token 进去。你可能会想：

为何每次翻译开头都是 start token，Transformer 还能产生不一样且正确的结果？

答案也很简单，因为 Decoder 可以透过「关注」 Encoder 处理完不同英文句子的输出来获得语义信息，了解它在当下该生成什么中文字作为第一个输出。

现在让我们定义一个 evaluate函数实现上述逻辑。此函数的输入是一个完全没有经过处理的英文句子（以字串表示），输出则是一个索引序列，里头的每个索引就代表着 Transformer 预测的中文字。

让我们实际看看 evaluate 函数：

# 给定一个英文句子，输出预测的中文索引数字序列以及注意权重 dict
def evaluate(inp_sentence):
  
  # 準備英文句子前後會加上的 <start>, <end>
  start_token = [subword_encoder_en.vocab_size]
  end_token = [subword_encoder_en.vocab_size + 1]
  
  # inp_sentence 是字串，我們用 Subword Tokenizer 將其變成子詞的索引序列
  # 並在前後加上 BOS / EOS
  inp_sentence = start_token + subword_encoder_en.encode(inp_sentence) + end_token
  encoder_input = tf.expand_dims(inp_sentence, 0)
  
  # 跟我們在影片裡看到的一樣，Decoder 在第一個時間點吃進去的輸入
  # 是一個只包含一個中文 <start> token 的序列
  decoder_input = [subword_encoder_zh.vocab_size]
  output = tf.expand_dims(decoder_input, 0)  # 增加 batch 維度
  
  # auto-regressive，一次生成一個中文字並將預測加到輸入再度餵進 Transformer
  for i in range(MAX_LENGTH):
    # 每多一個生成的字就得產生新的遮罩
    enc_padding_mask, combined_mask, dec_padding_mask = create_masks(
        encoder_input, output)
  
    # predictions.shape == (batch_size, seq_len, vocab_size)
    predictions, attention_weights = transformer(encoder_input, 
                                                 output,
                                                 False,
                                                 enc_padding_mask,
                                                 combined_mask,
                                                 dec_padding_mask)
    

    # 將序列中最後一個 distribution 取出，並將裡頭值最大的當作模型最新的預測字
    predictions = predictions[: , -1:, :]  # (batch_size, 1, vocab_size)

    predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
    
    # 遇到 <end> token 就停止回傳，代表模型已經產生完結果
    if tf.equal(predicted_id, subword_encoder_zh.vocab_size + 1):
      return tf.squeeze(output, axis=0), attention_weights
    
    #將 Transformer 新預測的中文索引加到輸出序列中，讓 Decoder 可以在產生
    # 下個中文字的時候關注到最新的 `predicted_id`
    output = tf.concat([output, predicted_id], axis=-1)

  # 將 batch 的維度去掉後回傳預測的中文索引序列
  return tf.squeeze(output, axis=0), attention_weights

我知道这章代码很多很长，但搭配注解后你会发现它们实际上都不难，而且这也是你看这篇文章的主要目的：实际了解 Transformer 是怎么做英中翻译的。你不想只是纸上谈兵，对吧？

有了 evaluate 函数，要透过 Transformer 做翻译非常容易：

# 要被翻译的英文句子
sentence = "China, India, and others have enjoyed continuing economic growth."

# 取得预测的中文索引序列
predicted_seq, _ = evaluate(sentence)

# 过滤掉 <start> & <end> tokens 并用中文的 subword tokenizer 帮我们将索引序列还原回中文句子
target_vocab_size = subword_encoder_zh.vocab_size
predicted_seq_without_bos_eos = [idx for idx in predicted_seq if idx < target_vocab_size]
predicted_sentence = subword_encoder_zh.decode(predicted_seq_without_bos_eos)

print("sentence:", sentence)
print("-" * 20)
print("predicted_seq:", predicted_seq)
print("-" * 20)
print("predicted_sentence:", predicted_sentence)

sentence: China, India, and others have enjoyed continuing economic growth.
--------------------
predicted_seq: tf.Tensor(
[4205   16    4   36  378  100    8   35   32    4   33  111  945  189
   22   49  105   83    3], shape=(19,), dtype=int32)
--------------------
predicted_sentence: 中国、印度和其他国家都享受经济增长。

考虑到这个 Transformer 不算巨大（约 400 万个参数），且模型训练时用的数据集不大的情况下，我们达到相当不错的结果，你说是吧？在这个例子里头该翻的词汇都翻了出来，句子本身也还算自然。

transformer.summary()

Model: "transformer_1"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
encoder_2 (Encoder)          multiple                  1831808   
_________________________________________________________________
decoder_2 (Decoder)          multiple                  1596800   
_________________________________________________________________
dense_137 (Dense)            multiple                  542703    
=================================================================
Total params: 3,971,311
Trainable params: 3,971,311
Non-trainable params: 0
_________________________________________________________________

12. 可视化注意权重

除了其运算高度并行以及表现不错以外，Transformer 另外一个优点在于我们可以透过可视化注意权重（attention weights）来了解模型实际在生成序列的时候放「注意力」在哪里。别忘记我们当初在 Decoder layers 做完 multi-head attention 之后都将注意权重输出。现在正是它们派上用场的时候了。

先让我们看看有什么注意权重可以拿来可视化：

predicted_seq, attention_weights = evaluate(sentence)

# 在这边我们自动选择最后一个 Decoder layer 的 MHA 2，也就是 Decoder 关注 Encoder 的 MHA
layer_name = f"decoder_layer{num_layers}_block2"

print("sentence:", sentence)
print("-" * 20)
print("predicted_seq:", predicted_seq)
print("-" * 20)
print("attention_weights.keys():")
for layer_name, attn in attention_weights.items():
  print(f"{layer_name}.shape: {attn.shape}")
print("-" * 20)
print("layer_name:", layer_name)

sentence: China, India, and others have enjoyed continuing economic growth.
--------------------
predicted_seq: tf.Tensor(
[4205   16    4   36  378  100    8   35   32    4   33  111  945  189
   22   49  105   83    3], shape=(19,), dtype=int32)
--------------------
attention_weights.keys():
decoder_layer1_block1.shape: (1, 8, 19, 19)
decoder_layer1_block2.shape: (1, 8, 19, 15)
decoder_layer2_block1.shape: (1, 8, 19, 19)
decoder_layer2_block2.shape: (1, 8, 19, 15)
decoder_layer3_block1.shape: (1, 8, 19, 19)
decoder_layer3_block2.shape: (1, 8, 19, 15)
decoder_layer4_block1.shape: (1, 8, 19, 19)
decoder_layer4_block2.shape: (1, 8, 19, 15)
--------------------
layer_name: decoder_layer4_block2

block1 代表是Decoder layer 自己关注自己的MHA 1，因此倒数两个维度都跟中文序列长度相同；
block2 则是Decoder layer 用来关注Encoder 输出的MHA 2 ，在这边我们选择最后一个Decoder layer 的MHA 2来看Transformer 在生成中文序列时关注在英文句子的那些位置。

但首先，我们得要有一个绘图的函数才行：

import matplotlib as mpl
# 你可能会需要自行下载一个中文字体档案以让 matplotlib 正确显示中文
zhfont = mpl.font_manager.FontProperties(fname='tensorflow-datasets/SimHei.ttf')
plt.style.use("seaborn-whitegrid")

# 这个函数将英 -> 中翻译的注意权重视觉化（注意：我们将注意权重 transpose 以最佳化渲染结果
def plot_attention_weights(attention_weights, sentence, predicted_seq, layer_name, max_len_tar=None):
    
    fig = plt.figure(figsize=(17, 7))
    sentence = subword_encoder_en.encode(sentence)
  
    # 只显示中文序列前 `max_len_tar` 个字以避免画面太过壅挤
    if max_len_tar:
        predicted_seq = predicted_seq[:max_len_tar]
    else:
        max_len_tar = len(predicted_seq)
  
    # 将某一个特定 Decoder layer 里头的 MHA 1 或 MHA2 的注意权重拿出来并去掉 batch 维度
    attention_weights = tf.squeeze(attention_weights[layer_name], axis=0)  
    # (num_heads, tar_seq_len, inp_seq_len)
  
    # 将每个 head 的注意权重画出
    for head in range(attention_weights.shape[0]):
        ax = fig.add_subplot(2, 4, head + 1)

        # [注意]我为了将长度不短的英文子词显示在 y 轴，将注意权重做了 transpose
        attn_map = np.transpose(attention_weights[head][:max_len_tar, :])
        ax.matshow(attn_map, cmap='viridis')  # (inp_seq_len, tar_seq_len)
    
        fontdict = {"fontproperties": zhfont}
    
        ax.set_xticks(range(max(max_len_tar, len(predicted_seq))))
        ax.set_xlim(-0.5, max_len_tar -1.5)
    
        ax.set_yticks(range(len(sentence) + 2))
        ax.set_xticklabels([subword_encoder_zh.decode([i]) for i in predicted_seq 
                            if i < subword_encoder_zh.vocab_size], 
                           fontdict=fontdict, fontsize=18)    
    
        ax.set_yticklabels(
            ['<start>'] + [subword_encoder_en.decode([i]) for i in sentence] + ['<end>'], 
            fontdict=fontdict)
    
        ax.set_xlabel('Head {}'.format(head + 1))
        ax.tick_params(axis="x", labelsize=12)
        ax.tick_params(axis="y", labelsize=12)
        
      
    plt.tight_layout()
    plt.show()
    plt.close(fig)

这个函数不难，且里头不少是调整图片的细节设定因此我将它留给你自行参考。

比较值得注意的是因为我们在这篇文章是做英文（来源）到中文（目标）的翻译，注意权重的 shape 为：

(batch_size, num_heads, zh_seq_len, en_seq_len)

如果你直接把注意权重绘出的话 y 轴就会是每个中文字，而 x 轴则会是每个英文子词。而英文子词绘在 x 轴太占空间，我将每个注意权重都做 transpose 并呈现结果，这点你得注意一下。

让我们马上画出刚刚翻译的注意权重看看：

plot_attention_weights(attention_weights, sentence, 
                       predicted_seq, layer_name, max_len_tar=18)

注意力权重可视化.png

尽管其运算机制十分错综复杂，阅读本文后 Transformer 对你来说不再是黑魔法，也不再是遥不可及的存在。如果你现在觉得「Transformer 也不过就这样嘛！」那就达成我写这篇文章的目的了。

自注意力机制以及Transformer 在推出之后就被非常广泛地使用并改进，但在我自己开始接触相关知识以后一直没有发现完整的繁中教学，因此写了这篇当初的我殷殷期盼的文章，也希望能帮助到更多人学习。

在进入结语之前，让我们看看文中的 Transformer 是怎么逐渐学会做好翻译的：

attention_weights_change_by_time.gif

13. 在你离开之前

这篇是当初在学习 Transformer 的我希望有人分享给自己的文章。

我相信人类之所以强大是因为集体知识：我们能透过书籍、影片以及语言将一个人脑中的知识与思想共享给其他人，让宝贵的知识能够「scale」，在更多人的脑袋中发光发热，创造更多价值。

我希望你有从本文中学到一点东西，并帮助我将本文的这些知识「scale」，把文章分享给更多有兴趣的人，并利用所学应用在一些你一直想要完成的任务上面。

最后一点提醒，就算Transformer比古早时代的方法好再多多终究也只是个工具，其最大价值不会超过于被你拿来应用的问题之上。就好像现在已有很多超越基本Transformer的翻译方法，但我们仍然持续在追寻更好的机器翻译系统。

工具会被淘汰，需求一直都在。