【从 0 开始学习 Transformer】下：Transformer 训练与评估

• 1. 前言
• 2. 创造原训练集的编码表示
  • 2.1. 数据下载与读取
  • 2.2. 创建子词分词器
  • 2.3. 数据处理
• 3. 损失函数设计
• 4. 优化器与学习率
• 5. 自回归原理
• 6. 训练
  • 6.1. 超参数
  • 6.2. 训练
    • 6.2.1. 创建遮挡
    • 6.2.2. 创建训练步骤及保存模型
    • 6.2.3. 开始训练
• 7. 评估

1. 前言

在上一篇文章中我们已经描述了 Transformer 的整个模型搭建过程，并逐层逐行地解释了其正向传播的原理和细节。接下来，我们将着手定义优化训练的方式，处理语料，并最终使用搭建好的 Transformer 实现一个由葡萄牙语翻译至英语的翻译器。

为了训练一个由葡萄牙语翻译至英语的翻译器，首先来观察如何处理数据从而能够正确地输入我们已经设计好的 Tranformer 模型：

class Transformer(tf.keras.Model):
...
    
  def call(self, inp, tar, training, enc_padding_mask, 
           look_ahead_mask, dec_padding_mask):

... 

只摘取模型的调用 call 部分，可以看出 Transformer 需要的输入：

• inp：输入序列，这里需要的是源语言（葡萄牙语）的编码表示。（嵌入表示将在编码器中完成）
• tar：目标序列，这里需要的是目标语言（英语）的编码表示。（嵌入表示将在编码器中完成）
• training：布尔量，规定模型是否可以训练。
• enc_padding_mask：编码器，填充遮挡。
• look_ahead_mask：前瞻遮挡。两个遮挡将在后面详细描述。
• dec_padding_mask：解码器，填充遮挡。

由此，我们知道，为了达成目的，我们需要完成以下几个步骤：

1. 创造原训练集（输入句子和目标句子）的嵌入表示
2. 为我们的 Transformer 设计优化器和损失函数
3. 根据情况创造填充遮挡
4. 为了实现自回归创建前瞻遮挡
5. 将数据输入进行训练
6. 最终对训练好的模型进行评估

2. 创造原训练集的编码表示

2.1. 数据下载与读取

参考 Tensorflow 的官方教程，我们同样使用 TFDS 来进行数据的下载和载入。（应首先在本机环境或虚拟环境中安装 tensorflow_datasets 模块。

import tensorflow_datasets as tfds

examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en', with_info=True,
                               as_supervised=True)
train_examples, val_examples = examples['train'], examples['validation']

第一行代码会访问用户目录下（Windows和Unix系系统各有不同，请参考官方文档）是否已经下载好了葡萄牙翻译至英文翻译器所需的数据集，如果不存在，则会自动下载。第二行，则将其自动转换为训练集合和测试集合两个 tf.data.Dataset 实例。

2.2. 创建子词分词器

tfds 独立于 Tensorflow，是专门用来管理和下载一些成熟的数据集的Python库。但其中有很多我认为通用性很强的函数。比如子词分词器：

tokenizer_en = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (en.numpy() for pt, en in train_examples), target_vocab_size=2**13)

tokenizer_pt = tfds.features.text.SubwordTextEncoder.build_from_corpus(
    (pt.numpy() for pt, en in train_examples), target_vocab_size=2**13)

上方代码分别创建了两个子词分词器，分别读取了训练集合中的全部英文和葡萄牙文，并基于这些大段的文字形成了子词分词器。

子词分词器的作用是将输入句子中的每一个单词编码为一个独一无二的数字，如果出现了子词分词器不能识别的新单词，那么就将其打散成多个可以识别的子词来编码成数字。

同样的，分词器也可以将用数字表示的句子重新转换回原有的句子。

sample_string = 'Transformer is awesome.'

tokenized_string = tokenizer_en.encode(sample_string)
print ('Tokenized string is {}'.format(tokenized_string))

# Tokenized string is [7915, 1248, 7946, 7194, 13, 2799, 7877]

original_string = tokenizer_en.decode(tokenized_string)
print ('The original string: {}'.format(original_string))

# The original string: Transformer is awesome.

assert original_string == sample_string

# 分词器转换回的句子和原始句子一定是相同的。

2.3. 数据处理

为了方便后期使用，编写一个将编码后句子加上开始标记和结束标记。利用 tf.data.Dataset 的 map 功能来批量完成这一任务。首先需要定义一个函数：

def encode(lang1, lang2):
  lang1 = [tokenizer_pt.vocab_size] + tokenizer_pt.encode(
      lang1.numpy()) + [tokenizer_pt.vocab_size+1]

  lang2 = [tokenizer_en.vocab_size] + tokenizer_en.encode(
      lang2.numpy()) + [tokenizer_en.vocab_size+1]
  
  return lang1, lang2

显然这里我们使用了原生 Python 编写这个函数，这样的函数是不能为 map 所用的。我们需要使用 tf.py_function 将这个函数转换为计算图。（此函数可以将原生 Python 编写的计算过程转换为 Tensorflow 流程控制的计算图，详情请参考 py_function)

def tf_encode(pt, en):
  return tf.py_function(encode, [pt, en], [tf.int64, tf.int64]) 
  # 第一个参数是包装的函数，第二个参数是输入的参数列表，第三个参数是输出的数据类型

于是，我们可以给训练数据集和验证数据集中每个句子加上开始标记和结束标记：

train_dataset = train_examples.map(tf_encode)
val_dataset = val_examples.map(tf_encode)

为了能够让这个模型较小，我们只使用句子短于 40 个单词的句子来作为输入数据。这里利用 tf.data.Dataset 的 filter 过滤器功能来快速筛选出需要的数据。

为了使用 filter，首先要定义一个过滤器函数。这是一个布尔函数，如果一条数据符合要求，则返回真，否则返回假。显然，对于葡萄牙句子翻译至英文句子数据集，我们要筛选出所有成对相同意思的句子，并且两条句子都短于 40 个单词（编码后并加上了开始和终结标记后的长度）。

def filter_max_length(x, y, max_length=MAX_LENGTH):
  return tf.logical_and(tf.size(x) <= max_length,
                        tf.size(y) <= max_length)

类似的，对数据集进行筛选：

train_dataset = train_dataset.filter(filter_max_length)

我们已经知道， 输入给 Transformer 的句子通常不会单句地输入，而是把句子叠成一批输入。将一批有长有短的句子叠成一批，需要将较短的句子补 0 使其长度匹配当前一批中最长的句子。

# 将数据集缓存到内存中以加快读取速度。
train_dataset = train_dataset.cache()
# shuffle 函数定义一个随机方式，首先定义一个缓存大小，取一部分数据放入缓存（BUFFER_SIZE大小），然后进行随机洗牌，最后从缓存中取。显然，若想实现全数据集的完美随机，需要让缓存的大小大于等于整个数据集。
# 首先将数据进行随机打散之后，对较短的数据进行填充。
train_dataset = train_dataset.shuffle(BUFFER_SIZE).padded_batch(
    BATCH_SIZE, padded_shapes=([-1], [-1]))
train_dataset = train_dataset.prefetch(tf.data.experimental.AUTOTUNE)

显然，验证集合也需要进行类似的处理操作（验证操作无需随机）。

val_dataset = val_dataset.filter(filter_max_length).padded_batch(
    BATCH_SIZE, padded_shapes=([-1], [-1]))

取出一个数据看一看：

pt_batch, en_batch = next(iter(val_dataset))
"""
pt_batch:
(,
en_batch:
 )
"""

3. 损失函数设计

损失函数的设计较为简单，需要考虑输出的句子和真正的目标句子是否为同一句子。只需要使用一个交叉熵函数。有一点需要注意，由上一章数据处理可以看出，数据中含有大量的填充（补0），这些填充不能作为真正的输入来考虑，因此在损失函数的计算中，需要将这些部分屏蔽掉。

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

def loss_function(real, pred):
# 对于mask，如果编码句子中出现了值为 0 的数据，则将其置 0
  mask = tf.math.logical_not(tf.math.equal(real, 0))
# 输出句子和真正的句子计算交叉熵
  loss_ = loss_object(real, pred)
# 将无效的交叉熵删除
  mask = tf.cast(mask, dtype=loss_.dtype)
  loss_ *= mask
# 返回平均值
  return tf.reduce_mean(loss_)

同时，定义两个指标用于展示训练过程中的模型变化：

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

4. 优化器与学习率

Transformer 使用 Adam 优化器，其 ${\beta }_1$ 为 0.9， ${\beta }_2$ 为0.98, $ϵ$ 为 $10^{-9}$。其学习率随着训练的进程变化：

其中，这个 warmup_step 设定为 4000。如此设计，学习率随着训练（Train Step）的变化就如下图所示

学习率的变化，我们通过继承 tf.keras.optimizers.schedules.LearningRateSchedule来实现。顾名思义，这个类会创建一个可序列化的学习率衰减（也可能增加）时间表：

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
  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): #　这个时间表被调用时，按照 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 = CustomSchedule(d_model)

optimizer = tf.keras.optimizers.Adam(learning_rate, beta_1=0.9, beta_2=0.98, 
                                     epsilon=1e-9)

5. 自回归原理

至今位置，我们已经拥有了 Transformer 的完整模型，数据输入和优化器。

但显然，Transformer 和 传统的 RNN 按时序依次读取输入和输出的训练方式“看起来”不同——它一次输入整个句子。而 encoder-decoder 架构是自回归的：通过上一步产生的符号和这一步的输入来预测这一步的输出。开始训练之前，需要了解 Transformer 是如何实现自回归的。

Tranformer 使用导师监督（teacher-forcing）法，即在预测过程中无论模型在当前时间步骤下预测出什么，teacher-forcing 方法都会将真实的输出传递到下一个时间步骤上。

当 transformer 预测每个词时，自注意力（self-attention）功能使它能够查看输入序列中前面的单词，从而更好地预测下一个单词。为了仅能让其查看输入序列中前面的单词，则需要前瞻遮挡来屏蔽后方的单词。

也就是说，若输入一个葡萄牙文句子，Tranformer 将第一次仅预测出英文句子的第一个单词，然后再次基础上依次预测第二个，第三个。

而训练过程也应该模拟这样的预测过程，每次仅增加一个目标序列的单词。

因此，我们将目标句子改写成两种：

原目标句子：sentence = "SOS A lion in the jungle is sleeping EOS"

改写为：

tar_inp = "SOS A lion in the jungle is sleeping"

tar_real = "A lion in the jungle is sleeping EOS"

（SOS 和 EOS 是开始标记和结束标记。）

真正输入给 Decoder 部分的是前者，配合前瞻遮挡它将模拟逐个单词产生的模型历史预测。而后者，则代表着模型当前步骤应该依次预测出的单词序列。很显然，他们应该仅仅只有一个单词的位移。

6. 训练

6.1. 超参数

Transformer 的基础模型使用的数值为：num_layers=6，d_model = 512，dff = 2048。

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

input_vocab_size = tokenizer_pt.vocab_size + 2
target_vocab_size = tokenizer_en.vocab_size + 2
dropout_rate = 0.1

6.2. 训练

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

6.2.1. 创建遮挡

• 首先要对输入数据（原始句子和目标句子）创建填充遮挡（填充了 0 的位置标记为 1，其余部分标记为 0，这里与损失函数的部分刚好相反）。
• 对于编码器解码器结构，当编码器预测后方的单词，只使用前方已经预测出的单词。为了实现这一效果，需要使用前瞻遮挡。

无论哪种遮挡，0 标记着保留的部分，1 标记着要遮挡的部分。

# 这里 inp 和 tar 都是来自第 2 章的编码后数据，形状显然为 (batch_size, len)
def create_masks(inp, tar):
  # 编码器填充遮挡，编码器自注意力时使用，在自注意力编码时排除掉没有含义的填充
  enc_padding_mask = create_padding_mask(inp)
  
  # 在解码器的第二个注意力模块使用。
  # 该填充遮挡用于遮挡编码器的输出，其输出输送给解码器使用，排除掉没有含义的填充
  dec_padding_mask = create_padding_mask(inp)
  
  # 在解码器的第一个注意力模块使用。
  # 遮挡（mask）解码器获取到的输入的后续标记（future tokens）。
  # 自然，填充的 padding 也不能忘记考虑，所以把两个遮挡合在一起两全其美
  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

6.2.2. 创建训练步骤及保存模型

为了保存模型，需要创建一个检查点管理器，在需要时使用此管理器来保存模型：

checkpoint_path = "./checkpoints/train"

ckpt = tf.train.Checkpoint(transformer=transformer,
                           optimizer=optimizer)

ckpt_manager = tf.train.CheckpointManager(ckpt, checkpoint_path, max_to_keep=5)

# 如果检查点存在，则恢复最新的检查点。
if ckpt_manager.latest_checkpoint:
  ckpt.restore(ckpt_manager.latest_checkpoint)
  print ('Latest checkpoint restored!!')

根据 5. 节的描述，将对目标序列进行调整和创建遮挡。最终实现训练过程。

在TF2.0中，由于使用了 eager excution 导致的性能下降，将使用@tf.function 装饰器将代码转换为传统的计算图提升性能。但这种转换并非完全智能，若没有良好的限制，则会因为输入 Tensor 的变化导致无法复用已有的计算图，导致冗余的转换。详情请参考Graph Execution 模式。

# 该 @tf.function 将追踪-编译 train_step 并将其转换为计算图，以便更快地执行。
# 该函数专用于参数张量的精确形状。为了避免由于可变序列长度或可变批次大小（最后一批次较小）导致的多次冗余转换
# 使用 input_signature 指定更多的通用形状。

train_step_signature = [
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
    tf.TensorSpec(shape=(None, None), dtype=tf.int64),
]

@tf.function(input_signature=train_step_signature)
def train_step(inp, tar):
  tar_inp = tar[:, :-1]
  tar_real = tar[:, 1:]
  
  enc_padding_mask, combined_mask, dec_padding_mask = create_masks(inp, tar_inp)
  
  with tf.GradientTape() as tape:
    predictions, _ = transformer(inp, tar_inp, 
                                 True, 
                                 enc_padding_mask, 
                                 combined_mask, 
                                 dec_padding_mask)
    loss = loss_function(tar_real, predictions)

  gradients = tape.gradient(loss, transformer.trainable_variables)    
  optimizer.apply_gradients(zip(gradients, transformer.trainable_variables))
  
  train_loss(loss)
  train_accuracy(tar_real, predictions)

6.2.3. 开始训练

EPOCHS = 20

for epoch in range(EPOCHS):
  start = time.time()
  
  train_loss.reset_states()
  train_accuracy.reset_states()
  
  # inp -> portuguese, tar -> english
  for (batch, (inp, tar)) in enumerate(train_dataset):
    train_step(inp, tar)
    
    if batch % 50 == 0:
      print ('Epoch {} Batch {} Loss {:.4f} Accuracy {:.4f}'.format(
          epoch + 1, batch, train_loss.result(), train_accuracy.result()))
      
  if (epoch + 1) % 5 == 0:
    ckpt_save_path = ckpt_manager.save()
    print ('Saving checkpoint for epoch {} at {}'.format(epoch+1,
                                                         ckpt_save_path))
    
  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))

效果如下：

Epoch 1 Batch 0 Loss 4.4721 Accuracy 0.0000
Epoch 1 Batch 50 Loss 4.2211 Accuracy 0.0076
Epoch 1 Batch 100 Loss 4.1943 Accuracy 0.0173
Epoch 1 Batch 150 Loss 4.1539 Accuracy 0.0205
Epoch 1 Batch 200 Loss 4.0675 Accuracy 0.0221
...

7. 评估

以下步骤用于评估：

• 用葡萄牙语分词器（tokenizer_pt）编码输入语句。此外，添加开始和结束标记，这样输入就与模型训练的内容相同。这是编码器输入。
• 解码器输入为 start token == tokenizer_en.vocab_size。
• 计算填充遮挡和前瞻遮挡。
• 解码器通过查看编码器输出和它自身的输出（自注意力）给出预测。
• 选择最后一个词并计算它的 argmax。将预测的词连接到解码器输入，然后传递给解码器。在这种方法中，解码器根据它预测的之前的词预测下一个。

评估函数：

def evaluate(inp_sentence):
  start_token = [tokenizer_pt.vocab_size]
  end_token = [tokenizer_pt.vocab_size + 1]
  
  # 输入语句是葡萄牙语，增加开始和结束标记
  inp_sentence = start_token + tokenizer_pt.encode(inp_sentence) + end_token
  encoder_input = tf.expand_dims(inp_sentence, 0)
  
  # 因为目标是英语，输入 transformer 的第一个词应该是
  # 英语的开始标记。
  decoder_input = [tokenizer_en.vocab_size]
  output = tf.expand_dims(decoder_input, 0)
    
  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)
    
    # 从 seq_len 维度选择最后一个词
    predictions = predictions[: ,-1:, :]  # (batch_size, 1, vocab_size)

    predicted_id = tf.cast(tf.argmax(predictions, axis=-1), tf.int32)
    
    # 如果 predicted_id 等于结束标记，就返回结果
    if predicted_id == tokenizer_en.vocab_size+1:
      return tf.squeeze(output, axis=0), attention_weights
    
    # 连接 predicted_id 与输出，作为解码器的输入传递到解码器。
    output = tf.concat([output, predicted_id], axis=-1)

  return tf.squeeze(output, axis=0), attention_weights

可视化注意力:

def plot_attention_weights(attention, sentence, result, layer):
  fig = plt.figure(figsize=(16, 8))
  
  sentence = tokenizer_pt.encode(sentence)
  
  attention = tf.squeeze(attention[layer], axis=0)
  
  for head in range(attention.shape[0]):
    ax = fig.add_subplot(2, 4, head+1)
    
    # 画出注意力权重
    ax.matshow(attention[head][:-1, :], cmap='viridis')

    fontdict = {'fontsize': 10}
    
    ax.set_xticks(range(len(sentence)+2))
    ax.set_yticks(range(len(result)))
    
    ax.set_ylim(len(result)-1.5, -0.5)
        
    ax.set_xticklabels(
        ['']+[tokenizer_pt.decode([i]) for i in sentence]+[''], 
        fontdict=fontdict, rotation=90)
    
    ax.set_yticklabels([tokenizer_en.decode([i]) for i in result 
                        if i < tokenizer_en.vocab_size], 
                       fontdict=fontdict)
    
    ax.set_xlabel('Head {}'.format(head+1))
  
  plt.tight_layout()
  plt.show()

单句子测试：

  def translate(sentence, plot=''):
    result, attention_weights = evaluate(sentence)
    
    predicted_sentence = tokenizer_en.decode([i for i in result 
                                              if i < tokenizer_en.vocab_size])  

    print('Input: {}'.format(sentence))
    print('Predicted translation: {}'.format(predicted_sentence))
    
    if plot:
      plot_attention_weights(attention_weights, sentence, result, plot)

  translate("este é um problema que temos que resolver.")
print ("Real translation: this is a problem we have to solve .")

效果：

Input: este é um problema que temos que resolver.
Predicted translation: this is a problem that we have to solve .... now .
Real translation: this is a problem we have to solve .