【深度学习】创建和训练Transformer神经网络模型，将葡萄牙语翻译成英语

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

前言

Transformer 是一种深度神经网络，它利用自注意力机制取代了 CNN 和 RNN 。自注意力机制让 Transformer 能够轻松地在输入序列之间传递信息。

正如Google AI 博客文章中所解释的那样：

图 1：将 Transformer 应用于机器翻译。

用于机器翻译的神经网络通常包含一个编码器，用于读取输入句子并生成其表示。然后，解码器在参考编码器生成的表示的同时逐字生成输出句子。Transformer 首先为每个单词生成初始表示或嵌入…然后，它使用自注意力机制聚合来自所有其他单词的信息，根据整个上下文为每个单词生成一个新的表示，由填充的球表示。然后对所有单词并行重复此步骤多次，依次生成新的表示。

Transformer 是一种序列到序列的编码器-解码器模型，类似于NMT with Attention 教程中的模型。单层 Transformer 需要多写一点代码，但与编码器-解码器 RNN 模型几乎相同。唯一的区别是 RNN 层被替换为自注意力层。本教程构建了一个 4 层 Transformer，它更大、更强大，但从根本上来说并不更复杂。

在训练模型后，您将能够输入葡萄牙语句子并返回英语翻译。
Transformer 为何如此重要

• Transformer 擅长对序列数据（例如自然语言）进行建模。
• 与循环神经网络 (RNN)不同，Transformers 是可并行化的。这使得它们在 GPU 和 TPU 等硬件上非常高效。主要原因是 Transformers 用注意力机制取代了循环，计算可以同时进行。层输出可以并行计算，而不是像 RNN 那样串联计算。
• 与RNN（如seq2seq，2014）或卷积神经网络 (CNN)（例如ByteNet ）不同，Transformer 能够捕获输入或输出序列中远距离位置之间的数据中的远距离或长距离上下文和依赖关系。因此，可以学习更长的连接。注意力机制允许每个位置在每一层访问整个输入，而在 RNN 和 CNN 中，信息需要经过许多处理步骤才能移动很长的距离，这使得学习变得更加困难。
• Transformers 不对数据中的时空关系做任何假设。这对于处理一组对象（例如星际争霸单位）非常理想。
  
  图 3：在英语到法语翻译上训练的 Transformer 的第 5 层到第 6 层中，单词“it”的编码器自注意力分布（八个注意力头之一）。

1. 安装

首先安装TensorFlow Datasets用于加载数据集，以及安装TensorFlow Text用于文本预处理：

# Install the most re version of TensorFlow to use the improved
# masking support for `tf.keras.layers.MultiHeadAttention`.
apt install --allow-change-held-packages libcudnn8=8.1.0.77-1+cuda11.2
pip uninstall -y -q tensorflow keras tensorflow-estimator tensorflow-text
pip install protobuf~=3.20.3
pip install -q tensorflow_datasets
pip install -q -U tensorflow-text tensorflow

导入必要的模块：

import logging
import time

import numpy as np
import matplotlib.pyplot as plt

import tensorflow_datasets as tfds
import tensorflow as tf

import tensorflow_text

2. 数据处理

下载数据集和子词标记器，然后将其全部包装在tf.data.Dataset训练中。

2.1 下载数据集

使用 TensorFlow Datasets 加载葡萄牙语-英语翻译数据集D Talks Open Translation Project。该数据集包含约 52,000 个训练样本、1,200 个验证样本和 1,800 个测试样本。

examples, metadata = tfds.load('ted_hrlr_translate/pt_to_en',
                               with_info=True,
                               as_supervised=True)

train_examples, val_examples = examples['train'], examples['validation']

TensorFlow Datasets 返回的对象tf.data.Dataset产生文本示例对：

for pt_examples, en_examples in train_examples.batch(3).take(1):
  print('> Examples in Portuguese:')
  for pt in pt_examples.numpy():
    print(pt.decode('utf-8'))
  print()

  print('> Examples in English:')
  for en in en_examples.numpy():
    print(en.decode('utf-8'))

2.2 设置标记器

现在您已经加载了数据集，您需要对文本进行标记，以便每个元素都表示为一个标记或标记 ID（数字表示形式）。

标记化是将文本分解为“标记”的过程。根据标记器的不同，这些标记可以表示句子片段、单词、子单词或字符。要了解有关标记化的更多信息，请访问此指南。

本教程使用子词标记器教程中内置的标记器。该教程针对此数据集text.BertTokenizer优化了两个对象（一个用于英语，一个用于葡萄牙语），并以 TensorFlow格式导出它们。saved_model

下载、提取并导入saved_model：

model_name = 'ted_hrlr_translate_pt_en_converter'
tf.keras.utils.get_file(
    f'{model_name}.zip',
    f'https://storage.googleapis.com/download.tensorflow.org/models/{model_name}.zip',
    cache_dir='.', cache_subdir='', extract=True
)

tokenizers = tf.saved_model.load(model_name)

包含tf.saved_model两个文本标记器，一个用于英语，一个用于葡萄牙语。两者都具有相同的方法：

[item for item in dir(tokenizers.en) if not item.startswith('_')]

该tokenize方法将一批字符串转换为一批填充的标记 ID。此方法在标记之前拆分标点符号、小写字母并对输入进行 Unicode 规范化。该标准化在此处不可见，因为输入数据已经标准化。

print('> This is a batch of strings:')
for en in en_examples.numpy():
  print(en.decode('utf-8'))

encoded = tokenizers.en.tokenize(en_examples)

print('> This is a padded-batch of token IDs:')
for row in encoded.to_list():
  print(row)

该detokenize方法尝试将这些令牌 ID 转换回人类可读的文本：

round_trip = tokenizers.en.detokenize(encoded)

print('> This is human-readable text:')
for line in round_trip.numpy():
  print(line.decode('utf-8'))

较低级别的lookup方法将 token-ID 转换为 token 文本：

print('> This is the text split into tokens:')
tokens = tokenizers.en.lookup(encoded)
tokens

输出演示了子词标记化的“子词”方面。

例如将单词分解searchability为search和##ability，将单词serendipity分解为s、##ere、##nd和##ip、##ity

请注意，标记化的文本包含[START]和[END]标记。

数据集中每个示例的标记分布如下：

lengths = []

for pt_examples, en_examples in train_examples.batch(1024):
  pt_tokens = tokenizers.pt.tokenize(pt_examples)
  lengths.append(pt_tokens.row_lengths())

  en_tokens = tokenizers.en.tokenize(en_examples)
  lengths.append(en_tokens.row_lengths())
  print('.', end='', flush=True)

all_lengths = np.concatenate(lengths)

plt.hist(all_lengths, np.linspace(0, 500, 101))
plt.ylim(plt.ylim())
max_length = max(all_lengths)
plt.plot([max_length, max_length], plt.ylim())
plt.title(f'Maximum tokens per example: {max_length}');

2.3 使用tf.data设置数据管道

以下函数将批量文本作为输入，并将其转换为适合训练的格式。

1. 它将它们标记为不规则的批次。
2. 它将每个修剪得不超过MAX_TOKENS。
3. 它将目标（英语）标记拆分为输入和标签。这些标记会移动一步，这样在每个输入位置上都有下label一个标记的 ID。
4. 它将RaggedTensors 转换为填充的密集Tensors。
5. 它返回一(inputs, labels)对。

MAX_TOKENS=128
def prepare_batch(pt, en):
    pt = tokenizers.pt.tokenize(pt)      # Output is ragged.
    pt = pt[:, :MAX_TOKENS]    # Trim to MAX_TOKENS.
    pt = pt.to_tensor()  # Convert to 0-padded dense Tensor

    en = tokenizers.en.tokenize(en)
    en = en[:, :(MAX_TOKENS+1)]
    en_inputs = en[:, :-1].to_tensor()  # Drop the [END] tokens
    en_labels = en[:, 1:].to_tensor()   # Drop the [START] tokens

    return (pt, en_inputs), en_labels

下面的函数将文本示例数据集转换为批量数据以供训练。

1. 它对文本进行标记，并过滤掉太长的序列。（之所以包含batch/，unbatch是因为标记器在处理大批量时效率更高）。
2. 该cache方法确保该工作只执行一次。
3. 然后shuffle，dense_to_ragged_batch随机化顺序并组装一批例子。
4. 最后prefetch将数据集与模型并行运行，以确保在需要时数据可用。

BUFFER_SIZE = 20000
BATCH_SIZE = 64

def make_batches(ds):
  return (
      ds
      .shuffle(BUFFER_SIZE)
      .batch(BATCH_SIZE)
      .map(prepare_batch, tf.data.AUTOTUNE)
      .prefetch(buffer_size=tf.data.AUTOTUNE))

3. 测试数据集

# Create training and validation set batches.
train_batches = make_batches(train_examples)
val_batches = make_batches(val_examples)

生成的tf.data.Dataset对象已设置为使用 Keras 进行训练。KerasModel.fit训练需要成对**(inputs, labels)的序列。inputs是标记化的葡萄牙语和英语序列对，(pt, en)是labels**相同的英语序列，但移位了 1。这种移位使得在每个位置输入en序列，label在下一个标记中。

下图：输入在底部，标签在顶部。

这与文本生成教程相同，只是这里有模型“条件化”的额外输入“上下文”（葡萄牙语序列）。

这种设置称为“教师强制”，因为无论模型在每个时间步的输出如何，它都会获得真实值作为下一个时间步的输入。这是一种简单而有效的文本生成模型训练方法。它之所以高效，是因为您不需要按顺序运行模型，不同序列位置的输出可以并行计算。

您可能以为input, output, 对只是序列Portuguese, English。给定葡萄牙语序列，模型将尝试生成英语序列。

可以用这种方式训练模型。您需要写出推理循环并将模型的输出传回输入。这种方式速度较慢（时间步骤不能并行运行），学习难度也较大（模型无法正确得出句子的结尾，除非它能正确得出句子的开头），但它可以提供更稳定的模型，因为模型必须在训练期间学会纠正自己的错误。

for (pt, en), en_labels in train_batches.take(1):
  break

print(pt.shape)
print(en.shape)
print(en_labels.shape)

和en相同en_labels，只是移动了 1：

print(en[0][:10])
print(en_labels[0][:10])

4. 定义组件

Transformer 内部有很多事情要做。需要记住的重要事项是：

1. 它遵循与带有编码器和解码器的标准序列到序列模型相同的一般模式。
2. 如果你一步一步地努力，一切都会变得有意义。

原始 Transformer 图 ------------------------ 4 层 Transformer 的表示

随着您学习本教程的进度，我们将解释这两个图中的每个组件。

4.1 嵌入和位置编码层

编码器和解码器的输入使用相同的嵌入和位置编码逻辑。

给定一个标记序列，输入标记（葡萄牙语）和目标标记（英语）都必须使用一个tf.keras.layers.Embedding层转换为向量。

整个模型中使用的注意层将其输入视为一组无序的向量。由于模型不包含任何循环层或卷积层。它需要某种方法来识别词序，否则它会将输入序列视为一个词袋，例如how are you，，，，等等，是无法区分的。how you areyou how are

Transformer 为嵌入向量添加了“位置编码”。它使用一组不同频率的正弦和余弦（跨序列）。根据定义，附近的元素将具有相似的位置编码。

原始论文采用以下公式来计算位置编码：

注意：下面的代码实现了它，但不是交错正弦和余弦，而是简单地连接正弦和余弦向量。像这样排列通道在功能上是等效的，而且更容易实现，并在下面的图中显示。

def positional_encoding(length, depth):
  depth = depth/2

  positions = np.arange(length)[:, np.newaxis]     # (seq, 1)
  depths = np.arange(depth)[np.newaxis, :]/depth   # (1, depth)

  angle_rates = 1 / (10000**depths)         # (1, depth)
  angle_rads = positions * angle_rates      # (pos, depth)

  pos_encoding = np.concatenate(
      [np.sin(angle_rads), np.cos(angle_rads)],
      axis=-1) 

  return tf.cast(pos_encoding, dtype=tf.float32)

位置编码函数是一堆正弦和余弦，它们根据沿嵌入向量深度的位置以不同的频率振动。它们在位置轴上振动。

pos_encoding = positional_encoding(length=2048, depth=512)

# Check the shape.
print(pos_encoding.shape)

# Plot the dimensions.
plt.pcolormesh(pos_encoding.numpy().T, cmap='RdBu')
plt.ylabel('Depth')
plt.xlabel('Position')
plt.colorbar()
plt.show()

根据定义，这些向量与位置轴上的邻近向量很好地对齐。下面对位置编码向量进行归一化，并1000通过点积将位置向量与所有其他向量进行比较：

pos_encoding/=tf.norm(pos_encoding, axis=1, keepdims=True)
p = pos_encoding[1000]
dots = tf.einsum('pd,d -> p', pos_encoding, p)
plt.subplot(2,1,1)
plt.plot(dots)
plt.ylim([0,1])
plt.plot([950, 950, float('nan'), 1050, 1050],
         [0,1,float('nan'),0,1], color='k', label='Zoom')
plt.legend()
plt.subplot(2,1,2)
plt.plot(dots)
plt.xlim([950, 1050])
plt.ylim([0,1])

因此，使用它来创建一个PositionEmbedding查找标记的嵌入向量并添加位置向量的层：

class PositionalEmbedding(tf.keras.layers.Layer):
  def __init__(self, vocab_size, d_model):
    super().__init__()
    self.d_model = d_model
    self.embedding = tf.keras.layers.Embedding(vocab_size, d_model, mask_zero=True) 
    self.pos_encoding = positional_encoding(length=2048, depth=d_model)

  def compute_mask(self, *args, **kwargs):
    return self.embedding.compute_mask(*args, **kwargs)

  def call(self, x):
    length = tf.shape(x)[1]
    x = self.embedding(x)
    # This factor sets the relative scale of the embedding and positonal_encoding.
    x *= tf.math.sqrt(tf.cast(self.d_model, tf.float32))
    x = x + self.pos_encoding[tf.newaxis, :length, :]
    return x

embed_pt = PositionalEmbedding(vocab_size=tokenizers.pt.get_vocab_size().numpy(), d_model=512)
embed_en = PositionalEmbedding(vocab_size=tokenizers.en.get_vocab_size().numpy(), d_model=512)

pt_emb = embed_pt(pt)
en_emb = embed_en(en)

en_emb._keras_mask

4.2 添加并规范化

这些“Add & Norm”块分散在整个模型中。每个块都加入一个残差连接，并通过一个LayerNormalization层运行结果。

组织代码的最简单方法是围绕这些残差块。以下部分将为每个残差块定义自定义层类。

残差“添加和规范”块包含在内，以便提高训练效率。残差连接为梯度提供了直接路径（并确保向量由注意层更新而不是替换），而规范化则为输出保持合理的比例。

4.3 基础注意力层

整个模型都使用了注意力层。除了注意力的配置方式外，这些层都是相同的。每个层都包含一个layers.MultiHeadAttention、一个layers.LayerNormalization和一个layers.Add。

要实现这些注意层，请从仅包含组件层的简单基类开始。每个用例都将作为子类实现。这样编写的代码会多一点，但意图却很明确。

class BaseAttention(tf.keras.layers.Layer):
  def __init__(self, **kwargs):
    super().__init__()
    self.mha = tf.keras.layers.MultiHeadAttention(**kwargs)
    self.layernorm = tf.keras.layers.LayerNormalization()
    self.add = tf.keras.layers.Add()

在了解每种用法的具体细节之前，先快速回顾一下注意力的工作原理：

有两个输入：

1. 查询序列；正在处理的序列；执行关注的序列（底部）。
2. 上下文序列；正在关注的序列（左）。

输出具有与查询序列相同的形状。

常见的比较是，此操作就像字典查找。模糊的、可微的、矢量化的字典查找。
这是一个常规的 Python 字典，其中有 3 个键和 3 个值，通过单个查询传递。

d = {'color': 'blue', 'age': 22, 'type': 'pickup'}
result = d['color']

• squery就是您要查找的内容。
• key字典里有什么样的信息。
• 这value就是那个信息。
  当您在普通字典中查找 时query，字典会找到匹配的key，并返回其关联的value。query要么有匹配key，要么没有。您可以想象一个模糊字典，其中的键不必完全匹配。如果您d[“species”]在上面的字典中查找，也许您希望它返回，"pickup"因为这是查询的最佳匹配。

注意层会进行这样的模糊查找，但它不只是寻找最佳键。它values根据query每个匹配程度将 组合起来key。

它是如何工作的？在注意层中query、key和value分别是向量。注意层不会进行哈希查找，而是将query和key向量组合起来以确定它们的匹配程度，即“注意分数”。该层返回所有的平均数values，由“注意分数”加权。

查询序列的每个位置都提供一个query向量。上下文序列充当字典。上下文序列中的每个位置都提供一个key和value向量。输入向量不直接使用，该layers.MultiHeadAttention层包括layers.Dense在使用输入向量之前对其进行投影的层。

4.4 交叉注意层

Transformer 的核心是交叉注意力层。该层连接编码器和解码器。该层是模型中最直接使用注意力的层，它执行的任务与 NMT 中的注意力模块相同。

为了实现这一点，您在调用层时将目标序列x作为传递query，并将context序列作为传递：key/valuemha

class CrossAttention(BaseAttention):
  def call(self, x, context):
    attn_output, attn_scores = self.mha(
        query=x,
        key=context,
        value=context,
        return_attention_scores=True)

    # Cache the attention scores for plotting later.
    self.last_attn_scores = attn_scores

    x = self.add([x, attn_output])
    x = self.layernorm(x)

    return x

下面的漫画展示了信息如何流经这一层。柱状图表示上下文序列的加权和。

为了简单起见，未显示残余连接。

输出长度是序列的长度query，而不是上下文序列的长度key/value。

下图进一步简化了。无需绘制整个“注意力权重”矩阵。重点是每个query位置都可以看到上下文中的所有key/value对，但查询之间不会交换任何信息。
在样本输入上进行测试运行：

sample_ca = CrossAttention(num_heads=2, key_dim=512)

print(pt_emb.shape)
print(en_emb.shape)
print(sample_ca(en_emb, pt_emb).shape)

4.5 全局自注意力层

该层负责处理上下文序列，并沿其长度传播信息：

由于在生成翻译时上下文序列是固定的，因此信息可以双向流动。

在 Transformer 和自注意力出现之前，模型通常使用 RNN 或 CNN 来完成此任务：

RNN 和 CNN 都有其局限性。

1. RNN 允许信息在整个序列中流动，但它需要经过许多处理步骤才能到达那里（限制梯度流）。这些 RNN 步骤必须按顺序运行，因此 RNN 不太能够利用现代并行设备。
2. 在 CNN 中，每个位置都可以并行处理，但它仅提供有限的接受场。接受场仅随 CNN 层数线性增长，您需要堆叠多个卷积层才能在序列中传输信息（Wavenet通过使用扩张卷积来减少此问题）。

另一方面，全局自注意层允许每个序列元素直接访问每个其他序列元素，只需少量操作，并且所有输出都可以并行计算。
要实现这一层，您只需将目标序列x作为query和value参数传递给该mha层：

class GlobalSelfAttention(BaseAttention):
  def call(self, x):
    attn_output = self.mha(
        query=x,
        value=x,
        key=x)
    x = self.add([x, attn_output])
    x = self.layernorm(x)
    return x

sample_gsa = GlobalSelfAttention(num_heads=2, key_dim=512)

print(pt_emb.shape)
print(sample_gsa(pt_emb).shape)

坚持与以前相同的风格，你可以像这样绘制：

再次，为了清楚起见，省略了残差连接。
像这样绘制更加紧凑，并且同样准确：

4.6 因果自注意力层

对于输出序列，该层的作用与全局自注意层类似：

这需要与编码器的全局自注意层以不同的方式处理。

与文本生成教程和带注意力机制的 NMT教程一样，Transformers 是一种“自回归”模型：它们一次生成一个标记的文本，并将该输出反馈给输入。为了提高效率，这些模型确保每个序列元素的输出仅取决于前一个序列元素；这些模型是“因果”的。

单向 RNN 从定义上讲是因果的。要进行因果卷积，您只需填充输入并移动输出，使其正确对齐（使用layers.Conv1D(padding='causal')）。

因果模型在两个方面是有效的：

1. 在训练中，它允许您仅执行一次模型即可计算输出序列中每个位置的损失。
2. 在推理过程中，对于生成的每个新标记，您只需计算其输出，前一个序列元素的输出可以重复使用。
   • 对于 RNN，您只需要 RNN 状态来考虑之前的计算（传递return_state=True给 RNN 层的构造函数）。
   • 对于 CNN，你需要遵循Fast Wavenet的方法

要构建因果自注意力层，需要在计算注意力分数和求和注意力时使用适当的掩码value。
use_causal_mask = True如果你在调用它时传递给该MultiHeadAttention层，则会自动处理这个问题：

class CausalSelfAttention(BaseAttention):
  def call(self, x):
    attn_output = self.mha(
        query=x,
        value=x,
        key=x,
        use_causal_mask = True)
    x = self.add([x, attn_output])
    x = self.layernorm(x)
    return x

因果掩码确保每个位置只能访问它之前的位置：

再次，为了简单起见，省略了残差连接。

该层更紧凑的表示形式为：

测试该层：

sample_csa = CausalSelfAttention(num_heads=2, key_dim=512)

print(en_emb.shape)
print(sample_csa(en_emb).shape)

早期序列元素的输出不依赖于后续元素，因此在应用层之前或之后修剪元素都无关紧要：

out1 = sample_csa(embed_en(en[:, :3])) 
out2 = sample_csa(embed_en(en))[:, :3]

tf.reduce_max(abs(out1 - out2)).numpy()

注意：使用 Keras 蒙版时，无效位置的输出值定义不明确。因此，上述内容可能不适用于蒙版区域。

4.7 前馈网络

Transformer 还在编码器和解码器中包含了这个逐点前馈网络：

该网络由两个线性层（tf.keras.layers.Dense）组成，中间有一个 ReLU 激活，还有一个 dropout 层。与注意层一样，此处的代码还包括残差连接和规范化：

class FeedForward(tf.keras.layers.Layer):
  def __init__(self, d_model, dff, dropout_rate=0.1):
    super().__init__()
    self.seq = tf.keras.Sequential([
      tf.keras.layers.Dense(dff, activation='relu'),
      tf.keras.layers.Dense(d_model),
      tf.keras.layers.Dropout(dropout_rate)
    ])
    self.add = tf.keras.layers.Add()
    self.layer_norm = tf.keras.layers.LayerNormalization()

  def call(self, x):
    x = self.add([x, self.seq(x)])
    x = self.layer_norm(x) 
    return x

测试层，输出与输入形状相同：

sample_ffn = FeedForward(512, 2048)

print(en_emb.shape)
print(sample_ffn(en_emb).shape)

4.8 编码器层

编码器包含一组N编码器层。其中每个编码器层EncoderLayer包含一个GlobalSelfAttention和FeedForward层：

这是的定义EncoderLayer：

class EncoderLayer(tf.keras.layers.Layer):
  def __init__(self,*, d_model, num_heads, dff, dropout_rate=0.1):
    super().__init__()

    self.self_attention = GlobalSelfAttention(
        num_heads=num_heads,
        key_dim=d_model,
        dropout=dropout_rate)

    self.ffn = FeedForward(d_model, dff)

  def call(self, x):
    x = self.self_attention(x)
    x = self.ffn(x)
    return x

快速测试一下，输出将具有与输入相同的形状：

sample_encoder_layer = EncoderLayer(d_model=512, num_heads=8, dff=2048)

print(pt_emb.shape)
print(sample_encoder_layer(pt_emb).shape)

4.9 编码器

接下来构建编码器。

编码器由以下部分组成：

• PositionalEmbedding输入处的一层。
• 一层层的堆叠EncoderLayer。

class Encoder(tf.keras.layers.Layer):
  def __init__(self, *, num_layers, d_model, num_heads,
               dff, vocab_size, dropout_rate=0.1):
    super().__init__()

    self.d_model = d_model
    self.num_layers = num_layers

    self.pos_embedding = PositionalEmbedding(
        vocab_size=vocab_size, d_model=d_model)

    self.enc_layers = [
        EncoderLayer(d_model=d_model,
                     num_heads=num_heads,
                     dff=dff,
                     dropout_rate=dropout_rate)
        for _ in range(num_layers)]
    self.dropout = tf.keras.layers.Dropout(dropout_rate)

  def call(self, x):
    # `x` is token-IDs shape: (batch, seq_len)
    x = self.pos_embedding(x)  # Shape `(batch_size, seq_len, d_model)`.

    # Add dropout.
    x = self.dropout(x)

    for i in range(self.num_layers):
      x = self.enc_layers[i](x)

    return x  # Shape `(batch_size, seq_len, d_model)`.

测试编码器：

# Instantiate the encoder.
sample_encoder = Encoder(num_layers=4,
                         d_model=512,
                         num_heads=8,
                         dff=2048,
                         vocab_size=8500)

sample_encoder_output = sample_encoder(pt, training=False)

# Print the shape.
print(pt.shape)
print(sample_encoder_output.shape)  # Shape `(batch_size, input_seq_len, d_model)`.

4.10 解码器层

解码器的堆栈稍微复杂一些，每个解码器DecoderLayer包含一个CausalSelfAttention、一个CrossAttention和一个FeedForward层：

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

    self.causal_self_attention = CausalSelfAttention(
        num_heads=num_heads,
        key_dim=d_model,
        dropout=dropout_rate)

    self.cross_attention = CrossAttention(
        num_heads=num_heads,
        key_dim=d_model,
        dropout=dropout_rate)

    self.ffn = FeedForward(d_model, dff)

  def call(self, x, context):
    x = self.causal_self_attention(x=x)
    x = self.cross_attention(x=x, context=context)

    # Cache the last attention scores for plotting later
    self.last_attn_scores = self.cross_attention.last_attn_scores

    x = self.ffn(x)  # Shape `(batch_size, seq_len, d_model)`.
    return x

测试解码器层：

sample_decoder_layer = DecoderLayer(d_model=512, num_heads=8, dff=2048)

sample_decoder_layer_output = sample_decoder_layer(
    x=en_emb, context=pt_emb)

print(en_emb.shape)
print(pt_emb.shape)
print(sample_decoder_layer_output.shape)  # `(batch_size, seq_len, d_model)`

4.11 解码器

与 类似Encoder，Decoder由PositionalEmbedding和 堆栈组成DecoderLayer：

通过扩展来定义解码器tf.keras.layers.Layer：

class Decoder(tf.keras.layers.Layer):
  def __init__(self, *, num_layers, d_model, num_heads, dff, vocab_size,
               dropout_rate=0.1):
    super(Decoder, self).__init__()

    self.d_model = d_model
    self.num_layers = num_layers

    self.pos_embedding = PositionalEmbedding(vocab_size=vocab_size,
                                             d_model=d_model)
    self.dropout = tf.keras.layers.Dropout(dropout_rate)
    self.dec_layers = [
        DecoderLayer(d_model=d_model, num_heads=num_heads,
                     dff=dff, dropout_rate=dropout_rate)
        for _ in range(num_layers)]

    self.last_attn_scores = None

  def call(self, x, context):
    # `x` is token-IDs shape (batch, target_seq_len)
    x = self.pos_embedding(x)  # (batch_size, target_seq_len, d_model)

    x = self.dropout(x)

    for i in range(self.num_layers):
      x  = self.dec_layers[i](x, context)

    self.last_attn_scores = self.dec_layers[-1].last_attn_scores

    # The shape of x is (batch_size, target_seq_len, d_model).
    return x

测试解码器：

# Instantiate the decoder.
sample_decoder = Decoder(num_layers=4,
                         d_model=512,
                         num_heads=8,
                         dff=2048,
                         vocab_size=8000)

output = sample_decoder(
    x=en,
    context=pt_emb)

# Print the shapes.
print(en.shape)
print(pt_emb.shape)
print(output.shape)

sample_decoder.last_attn_scores.shape  # (batch, heads, target_seq, input_seq)

创建了 Transformer 编码器和解码器后，就该构建 Transformer 模型并进行训练了。

5. 变压器

现在你有Encoder和Decoder。要完成模型Transformer，你需要将它们放在一起并添加最终的线性（Dense）层，该层将每个位置的结果向量转换为输出标记概率。

解码器的输出是该最终线性层的输入。

和Transformer中各有一层的模型看起来几乎与RNN+attention 教程中的模型完全相同。多层 Transformer 具有更多层，但本质上做的事情相同。EncoderDecoder

Transformer通过扩展来创建tf.keras.Model：

class Transformer(tf.keras.Model):
  def __init__(self, *, num_layers, d_model, num_heads, dff,
               input_vocab_size, target_vocab_size, dropout_rate=0.1):
    super().__init__()
    self.encoder = Encoder(num_layers=num_layers, d_model=d_model,
                           num_heads=num_heads, dff=dff,
                           vocab_size=input_vocab_size,
                           dropout_rate=dropout_rate)

    self.decoder = Decoder(num_layers=num_layers, d_model=d_model,
                           num_heads=num_heads, dff=dff,
                           vocab_size=target_vocab_size,
                           dropout_rate=dropout_rate)

    self.final_layer = tf.keras.layers.Dense(target_vocab_size)

  def call(self, inputs):
    # To use a Keras model with `.fit` you must pass all your inputs in the
    # first argument.
    context, x  = inputs

    context = self.encoder(context)  # (batch_size, context_len, d_model)

    x = self.decoder(x, context)  # (batch_size, target_len, d_model)

    # Final linear layer output.
    logits = self.final_layer(x)  # (batch_size, target_len, target_vocab_size)

    try:
      # Drop the keras mask, so it doesn't scale the losses/metrics.
      # b/250038731
      del logits._keras_mask
    except AttributeError:
      pass

    # Return the final output and the attention weights.
    return logits

5.1 超参数

为了使这个例子保持较小并且相对较快，层数（num_layers）、嵌入的维数（d_model）和层的内部维数FeedForward（dff）都已经减少了。

原始 Transformer 论文中描述的基础模型使用了num_layers=6、d_model=512和dff=2048。

自注意力头的数量保持不变（num_heads=8）。

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

5.2 尝试一下

实例化Transformer模型：

transformer = Transformer(
    num_layers=num_layers,
    d_model=d_model,
    num_heads=num_heads,
    dff=dff,
    input_vocab_size=tokenizers.pt.get_vocab_size().numpy(),
    target_vocab_size=tokenizers.en.get_vocab_size().numpy(),
    dropout_rate=dropout_rate)

测试一下：

output = transformer((pt, en))

print(en.shape)
print(pt.shape)
print(output.shape)

attn_scores = transformer.decoder.dec_layers[-1].last_attn_scores
print(attn_scores.shape)  # (batch, heads, target_seq, input_seq)

打印模型摘要：

transformer.summary()

6. 训练模型

现在是时候准备模型并开始训练它了。

6.1 设置优化器

根据原始 Transformer论文中的公式，使用带有自定义学习率调度器的 Adam 优化器。

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
  def __init__(self, d_model, warmup_steps=4000):
    super().__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 = tf.cast(step, dtype=tf.float32)
    arg1 = tf.math.rsqrt(step)
    arg2 = step * (self.warmup_steps ** -1.5)

    return tf.math.rsqrt(self.d_model) * tf.math.minimum(arg1, arg2)

实例化优化器（在此示例中为tf.keras.optimizers.Adam）：

learning_rate = CustomSchedule(d_model)

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

测试自定义学习率调度程序：

plt.plot(learning_rate(tf.range(40000, dtype=tf.float32)))
plt.ylabel('Learning Rate')
plt.xlabel('Train Step')

6.2 设置损失和指标

由于目标序列是填充的，因此在计算损失时应用填充掩码非常重要。使用交叉熵损失函数 ( tf.keras.losses.SparseCategoricalCrossentropy)：

def masked_loss(label, pred):
  mask = label != 0
  loss_object = tf.keras.losses.SparseCategoricalCrossentropy(
    from_logits=True, reduction='none')
  loss = loss_object(label, pred)

  mask = tf.cast(mask, dtype=loss.dtype)
  loss *= mask

  loss = tf.reduce_sum(loss)/tf.reduce_sum(mask)
  return loss


def masked_accuracy(label, pred):
  pred = tf.argmax(pred, axis=2)
  label = tf.cast(label, pred.dtype)
  match = label == pred

  mask = label != 0

  match = match & mask

  match = tf.cast(match, dtype=tf.float32)
  mask = tf.cast(mask, dtype=tf.float32)
  return tf.reduce_sum(match)/tf.reduce_sum(mask)

6.3 训练模型

所有组件准备就绪后，使用 配置训练过程model.compile，然后使用 运行它model.fit：

transformer.compile(
    loss=masked_loss,
    optimizer=optimizer,
    metrics=[masked_accuracy])

transformer.fit(train_batches,
                epochs=20,
                validation_data=val_batches)

7. 运行推理

现在，您可以通过执行翻译来测试模型。以下步骤用于推理：

• 使用葡萄牙语标记器 ( ) 对输入句子进行编码tokenizers.pt。这是编码器输入。
• 解码器输入被初始化为[START]令牌。
• 计算填充掩码和前瞻掩码。
• 然后decoder通过查看encoder output和其自身的输出（自我注意力）来输出预测。
• 将预测的标记连接到解码器输入并将其传递给解码器。
• 在这种方法中，解码器根据其预测的前一个标记来预测下一个标记。

注意：该模型针对高效训练进行了优化，并同时对输出中的每个标记进行下一个标记预测。这在推理过程中是多余的，并且只使用最后一个预测。如果在推理模式下运行时仅计算最后一个预测，则可以使该模型的推理效率更高（）。training=False

Translator通过子类化来定义类tf.Module：

class Translator(tf.Module):
  def __init__(self, tokenizers, transformer):
    self.tokenizers = tokenizers
    self.transformer = transformer

  def __call__(self, sentence, max_length=MAX_TOKENS):
    # The input sentence is Portuguese, hence adding the `[START]` and `[END]` tokens.
    assert isinstance(sentence, tf.Tensor)
    if len(sentence.shape) == 0:
      sentence = sentence[tf.newaxis]

    sentence = self.tokenizers.pt.tokenize(sentence).to_tensor()

    encoder_input = sentence

    # As the output language is English, initialize the output with the
    # English `[START]` token.
    start_end = self.tokenizers.en.tokenize([''])[0]
    start = start_end[0][tf.newaxis]
    end = start_end[1][tf.newaxis]

    # `tf.TensorArray` is required here (instead of a Python list), so that the
    # dynamic-loop can be traced by `tf.function`.
    output_array = tf.TensorArray(dtype=tf.int64, size=0, dynamic_size=True)
    output_array = output_array.write(0, start)

    for i in tf.range(max_length):
      output = tf.transpose(output_array.stack())
      predictions = self.transformer([encoder_input, output], training=False)

      # Select the last token from the `seq_len` dimension.
      predictions = predictions[:, -1:, :]  # Shape `(batch_size, 1, vocab_size)`.

      predicted_id = tf.argmax(predictions, axis=-1)

      # Concatenate the `predicted_id` to the output which is given to the
      # decoder as its input.
      output_array = output_array.write(i+1, predicted_id[0])

      if predicted_id == end:
        break

    output = tf.transpose(output_array.stack())
    # The output shape is `(1, tokens)`.
    text = tokenizers.en.detokenize(output)[0]  # Shape: `()`.

    tokens = tokenizers.en.lookup(output)[0]

    # `tf.function` prevents us from using the attention_weights that were
    # calculated on the last iteration of the loop.
    # So, recalculate them outside the loop.
    self.transformer([encoder_input, output[:,:-1]], training=False)
    attention_weights = self.transformer.decoder.last_attn_scores

    return text, tokens, attention_weights

注意：此函数使用展开循环，而不是动态循环。MAX_TOKENS每次调用时都会生成。请参阅NMT with Attention教程，获取使用动态循环的示例实现，这样效率会高得多。

创建此类的一个实例Translator，并尝试几次：

translator = Translator(tokenizers, transformer)
def print_translation(sentence, tokens, ground_truth):
  print(f'{"Input:":15s}: {sentence}')
  print(f'{"Prediction":15s}: {tokens.numpy().decode("utf-8")}')
  print(f'{"Ground truth":15s}: {ground_truth}')

示例 1：

sentence = 'este é um problema que temos que resolver.'
ground_truth = 'this is a problem we have to solve .'

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)

示例 2：

sentence = 'os meus vizinhos ouviram sobre esta ideia.'
ground_truth = 'and my neighboring homes heard about this idea .'

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)

示例 3：

sentence = 'vou então muito rapidamente partilhar convosco algumas histórias de algumas coisas mágicas que aconteceram.'
ground_truth = "so i'll just share with you some stories very quickly of some magical things that have happened."

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)

8. 创建注意力图

您在上一节中创建的类Translator返回一个注意力热图词典，您可以使用它来可视化模型的内部工作。

例如：

sentence = 'este é o primeiro livro que eu fiz.'
ground_truth = "this is the first book i've ever done."

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)

创建一个函数，在生成标记时绘制注意力：

def plot_attention_head(in_tokens, translated_tokens, attention):
  # The model didn't generate `<START>` in the output. Skip it.
  translated_tokens = translated_tokens[1:]

  ax = plt.gca()
  ax.matshow(attention)
  ax.set_xticks(range(len(in_tokens)))
  ax.set_yticks(range(len(translated_tokens)))

  labels = [label.decode('utf-8') for label in in_tokens.numpy()]
  ax.set_xticklabels(
      labels, rotation=90)

  labels = [label.decode('utf-8') for label in translated_tokens.numpy()]
  ax.set_yticklabels(labels)

head = 0
# Shape: `(batch=1, num_heads, seq_len_q, seq_len_k)`.
attention_heads = tf.squeeze(attention_weights, 0)
attention = attention_heads[head]
attention.shape

这些是输入（葡萄牙语）标记：

in_tokens = tf.convert_to_tensor([sentence])
in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()
in_tokens = tokenizers.pt.lookup(in_tokens)[0]
in_tokens

这些是输出（英语翻译）标记：

translated_tokens
plot_attention_head(in_tokens, translated_tokens, attention)
def plot_attention_weights(sentence, translated_tokens, attention_heads):
  in_tokens = tf.convert_to_tensor([sentence])
  in_tokens = tokenizers.pt.tokenize(in_tokens).to_tensor()
  in_tokens = tokenizers.pt.lookup(in_tokens)[0]

  fig = plt.figure(figsize=(16, 8))

  for h, head in enumerate(attention_heads):
    ax = fig.add_subplot(2, 4, h+1)

    plot_attention_head(in_tokens, translated_tokens, head)

    ax.set_xlabel(f'Head {h+1}')

  plt.tight_layout()
  plt.show()
  
plot_attention_weights(sentence,
                       translated_tokens,
                       attention_weights[0])

该模型可以处理不熟悉的单词。 和’triceratops’都不’encyclopédia’在输入数据集中，即使没有共享词汇表，模型也会尝试音译它们。例如：

sentence = 'Eu li sobre triceratops na enciclopédia.'
ground_truth = 'I read about triceratops in the encyclopedia.'

translated_text, translated_tokens, attention_weights = translator(
    tf.constant(sentence))
print_translation(sentence, translated_text, ground_truth)

plot_attention_weights(sentence, translated_tokens, attention_weights[0])

9. 导出模型

您已测试模型并且推理正常。接下来，您可以将其导出为tf.saved_model。要了解如何以 SavedModel 格式保存和加载模型，请使用本指南。

通过使用该方法对子类进行子ExportTranslator类化来创建一个调用的类：tf.Moduletf.function__call__

class ExportTranslator(tf.Module):
  def __init__(self, translator):
    self.translator = translator

  @tf.function(input_signature=[tf.TensorSpec(shape=[], dtype=tf.string)])
  def __call__(self, sentence):
    (result,
     tokens,
     attention_weights) = self.translator(sentence, max_length=MAX_TOKENS)

    return result

上面tf.function只返回了输出句子。由于执行不严格，所以tf.function不会计算任何不必要的值。

包裹translator新创建的ExportTranslator：

translator = ExportTranslator(translator)

由于模型正在解码预测，因此tf.argmax预测是确定性的。原始模型和从其重新加载的模型SavedModel应该给出相同的预测：

translator('este é o primeiro livro que eu fiz.').numpy()
tf.saved_model.save(translator, export_dir='translator')
reloaded = tf.saved_model.load('translator')
reloaded('este é o primeiro livro que eu fiz.').numpy()

10. 总结

在本教程中，您学习了：

• Transformer 及其在机器学习中的重要性
• 注意力、自注意力和多头注意力
• 嵌入位置编码
• 原始 Transformer 的编码器-解码器架构
• 自我注意力的掩蔽
• 如何将所有内容整合在一起来翻译文本

这种架构的缺点是：

• 对于时间序列，时间步长的输出是根据整个历史记录计算得出的，而不是仅根据输入和当前隐藏状态计算得出的。这可能效率较低。
• 如果输入具有时间/空间关系，如文本或图像，则必须添加一些位置编码，否则模型将有效地看到一个词袋。

如果你想练习，你可以尝试很多事情。例如：

• 使用不同的数据集来训练 Transformer。
• 通过更改超参数，从原始论文创建“Base Transformer”或“Transformer XL”配置。
• 使用此处定义的层来创建BERT的实现
• 使用 Beam search 来获得更好的预测。

基于 Transformer 的模型种类繁多，其中许多模型在 2017 年版原始 Transformer 的基础上进行了改进，采用了编码器-解码器、仅编码器和仅解码器架构。