PyTorch Transformer 英中翻译超详细教程_问答-阿里云开发者社区

在这篇文章中，我将详细地介绍一个英中翻译 Transformer 的 PyTorch 实现。这篇文章会完整地展示一个深度学习项目的搭建过程，从数据集准备，到模型定义、训练。这篇文章不仅会讲解如何把 Transformer 的论文翻译成代码，还会讲清楚代码实现中的诸多细节。相信初学者能够从这篇文章中学到丰富的知识。

如果你只对 Transformer 感兴趣，而不想涉及过多的 NLP 内容，可以选读模型实现部分。

项目网址: https://github.com/SingleZombie/DL-Demos/pull/21

如果你对 Transformer 的论文不熟，欢迎阅读我之前的文章：Attention Is All You Need (Transformer) 论文精读。

任务与数据集准备

在开启一个深度学习项目之初，要把任务定义好。准确来说，我们要明白这个任务是在完成一个怎样的映射，并准备一个用于评测的数据集，定义好评价指标。

英中翻译，这个任务非常明确，就是把英文的句子翻译成中文。英中翻译的数据集应该包含若干个句子对，每个句子对由一句英文和它对应的中文翻译组成。

中英翻译的数据集不是很好找。有几个比较出名的数据集的链接已经失效了，还有些数据集需要注册与申请后才能获取。我在中文NLP语料库仓库(https://github.com/brightmart/nlp_chinese_corpus)找到了中英文平行语料 translation2019zh。该语料库由520万对中英文语料构成，训练集516万对，验证集3.9万对。用作训练和验证中英翻译模型是足够了。

机器翻译的评测指标叫做BLEU Score。如果模型输出的翻译和参考译文有越多相同的单词、连续2个相同单词、连续3个相同单词……，则得分越高。

PyTorch 提供了便捷的API，我们可以用一行代码算完BLEU Score。

>>> from torchtext.data.metrics import bleu_score
>>> candidate_corpus = [['My', 'full', 'pytorch', 'test'], ['Another', 'Sentence']]
>>> references_corpus = [[['My', 'full', 'pytorch', 'test'], ['Completely', 'Different']], [['No', 'Match']]]
>>> bleu_score(candidate_corpus, references_corpus)
    0.8408964276313782

数据集预处理

数据清洗

得到数据集后，下一步要做的是对数据集做处理，把原始数据转化成能够输入神经网络的张量。对于图片，预处理可能是裁剪、缩放，使所有图片都有一样的大小；对于文本，预处理可能是分词、填充。

在网盘上下载好 translation2019zh 数据集后，我们来一步一步清洗这个数据集。这个数据集只有两个文件translation2019zh_train.json, translation2019zh_valid.json，它们的结构如下：

{"english": <english>, "chinese": <chinese>}
{"english": <english>, "chinese": <chinese>}
{"english": <english>, "chinese": <chinese>}
...

这些json文件有点不合标准，每对句子由一行json格式的记录组成。english属性是英文句子，chinese属性是中文句子。比如：

{"english": "In Italy ...", "chinese": "在意大利 ..."}

因此，在读取数据时，我们可以用下面的代码提取每对句子。

import json

with open(json_path, 'r') as fp:
    for line in fp:
        line = json.loads(line)
        english, chinese = line['english'], line['chinese']

这个数据集有一点不干净，有一些句子对的中英文句子颠倒过来了。为此，我们要稍微处理一下，把这些句子对翻转过来。如果一个英文句子不全由 ASCII 组成，则它可能是一个被标错的中文句子。

# Correct mislabeled data
    if not english.isascii():
        english, chinese = chinese, english

经过这一步，我们只得到了中英文的字符文本。而在NLP中，大部分处理的最小单位都是符号（token）——对于英文来说，符号是单词、标点；对于中文来说，符号是词语、标点。我们还需要一个符号化的过程。

英文符号化非常方便，torchtext 提供了非常便捷的英文分词 API。

from torchtext.data import get_tokenizer

tokenizer = get_tokenizer('basic_english')
english = tokenizer(english)

而中文分词方面，我使用了jieba库。该库可以直接 pip 安装。

pip install jieba

分词的 API 是 jieba.cut。由于分词的结果中，相邻的词之间有空格，我一股脑地把所有空白符给过滤掉了。

import jieba
chinese = list(jieba.cut(chinese))
chinese = [x for x in chinese if x not in {' ', '\t'}]

经过这些处理后，每句话被转换成了中文词语或英文单词的数组。整个处理代码如下：

def read_file(json_path):
    english_sentences = []
    chinese_sentences = []
    tokenizer = get_tokenizer('basic_english')
    with open(json_path, 'r') as fp:
        for line in fp:
            line = json.loads(line)
            english, chinese = line['english'], line['chinese']
            # Correct mislabeled data
            if not english.isascii():
                english, chinese = chinese, english
            # Tokenize
            english = tokenizer(english)
            chinese = list(jieba.cut(chinese))
            chinese = [x for x in chinese if x not in {' ', '\t'}]
            english_sentences.append(english)
            chinese_sentences.append(chinese)
    return english_sentences, chinese_sentences

词语转序号

为了让计算机更方便地处理单词，我们还要把单词转换成序号。比如令apple为0号，banana为1号，则句子apple banana apple就转换成了0 1 0。

给每一个单词选一个标号，其实就是要建立一个词典。一般来说，我们可以利用他人的统计结果，挑选最常用的一些英文单词和中文词语构成词典。不过，现在我们已经有了一个庞大的中英语料库了，我们可以直接从这个语料库中挑选出最常见的词构成词典。

根据上一步处理得到的句子数组sentences，我们可以用下面的 Python 代码统计出最常见的一些词语，把它们和4个特殊字符, , , （句子开始字符、句子结束字符、频率太少没有被加入词典的词语、填充字符）一起构成词典。统计字符出现次数是通过 Python 的 Counter 类实现的。

from collections import Counter

def create_vocab(sentences, max_element=None):
    """Note that max_element includes special characters"""

    default_list = ['<sos>', '<eos>', '<unk>', '<pad>']

    char_set = Counter()
    for sentence in sentences:
        c_set = Counter(sentence)
        char_set.update(c_set)

    if max_element is None:
        return default_list + list(char_set.keys())
    else:
        max_element -= 4
        words_freq = char_set.most_common(max_element)
        # pair array to double array
        words, freq = zip(*words_freq)
        return default_list + list(words)

准备好了词典后，我还编写了两个工具函数sentence_to_tensor，tensor_to_sentence，它们可以用于字符串数组与序号数组的互相转换。测试这些代码的脚本及其输出如下：

# Dataset.py

def main():
    en_sens, zh_sens = read_file(
        'data/translation2019zh/translation2019zh_valid.json')
    print(*en_sens[0:3])
    print(*zh_sens[0:3])
    en_vocab = create_vocab(en_sens, 10000)
    zh_vocab = create_vocab(zh_sens, 30000)
    print(list(en_vocab)[0:10])
    print(list(zh_vocab)[0:10])

    en_tensors = sentence_to_tensor(en_sens, en_vocab)
    zh_tensors = sentence_to_tensor(zh_sens, zh_vocab)

    print(tensor_to_sentence(en_tensors[0], en_vocab, True))
    print(tensor_to_sentence(zh_tensors[0], zh_vocab))

['slowly', 'and', 'not', 'without', 'struggle', ',', 'america', 'began', 'to', 'listen', '.'] ...]
['美国', '缓慢', '地', '开始', '倾听', '，', '但', '并非', '没有', '艰难曲折', '。'] ...]
['<sos>', '<eos>', '<unk>', '<pad>', 'the', '.', ',', 'of', 'and', 'to']
['<sos>', '<eos>', '<unk>', '<pad>', '的', '，', '。', '在', '了', '和']
slowly and not without struggle , america began to listen .
美国缓慢地开始倾听，但并非没有<unk>。

在这一步中，有一个重要的参数：词典的大小。显然，词典越大，能处理的词语越多，但训练速度也会越慢。由于这个项目只是一个用于学习的demo，我设置了比较小的词典大小。想提升整个模型的性能的话，调大词典大小是一个最快的方法。

生成 Dataloader

都说程序员是新时代的农民工，这非常有道理。因为，作为程序员，你免不了要写一些繁重、无聊的数据处理脚本。还好，写完这些无聊的预处理代码后，总算可以使用 PyTorch 的 API 写一些有趣的代码了。

把词语数组转换成序号句子数组后，我们要考虑怎么把序号句子数组输入给模型了。文本数据通常长短不一，为了一次性处理一个 batch 的数据，要把短的句子填充，使得一批句子长度相等。写 Dataloader 时最主要的工作就是填充并对齐句子。

先看一下Dataset的写法。上一步得到的序号句子数组可以塞进Dataset里。注意，每个句子的前后要加上表示句子开始和结束的特殊符号。

SOS_ID = 0
EOS_ID = 1
UNK_ID = 2
PAD_ID = 3

class TranslationDataset(Dataset):

    def __init__(self, en_tensor: np.ndarray, zh_tensor: np.ndarray):
        super().__init__()
        assert len(en_tensor) == len(zh_tensor)
        self.length = len(en_tensor)
        self.en_tensor = en_tensor
        self.zh_tensor = zh_tensor

    def __len__(self):
        return self.length

    def __getitem__(self, index):
        x = np.concatenate(([SOS_ID], self.en_tensor[index], [EOS_ID]))
        x = torch.from_numpy(x)
        y = np.concatenate(([SOS_ID], self.zh_tensor[index], [EOS_ID]))
        y = torch.from_numpy(y)
        return x, y

接下来看一下 DataLoader 的写法。在创建 Dataloader 时，最重要的是 collate_fn 的编写，这个函数决定了怎么把多条数据合成一个等长的 batch。

def get_dataloader(en_tensor: np.ndarray,
                   zh_tensor: np.ndarray,
                   batch_size=16):

    def collate_fn(batch):
        ...

    dataset = TranslationDataset(en_tensor, zh_tensor)
    dataloader = DataLoader(dataset,
                            batch_size=batch_size,
                            shuffle=True,
                            collate_fn=collate_fn)

    return dataloader

collate_fn 的输入是多个 dataset __getitem__ 的返回结果构成的数组。对于我们的 dataset 来说，collate_fn 的输入是 [(x1, y1), (x2, y2), ...] 。我们可以用 zip(*batch) 把二元组数组拆成两个数组 x, y 。

collate_fn 的输出就是将来 dataloader 的输出。PyTorch 提供了 pad_sequence 函数用来把一批数据填充至等长。

from torch.nn.utils.rnn import pad_sequence

def collate_fn(batch):
    x, y = zip(*batch)
    x_pad = pad_sequence(x, batch_first=True, padding_value=PAD_ID)
    y_pad = pad_sequence(y, batch_first=True, padding_value=PAD_ID)

    return x_pad, y_pad

Transformer 模型

准备好 DataLoader 后，接下来就要进入这个项目最重要的部分——Transformer 模型实现了。我将按照代码的执行顺序，从前往后，自底向上地介绍 Transformer 的各个模块：Positional Encoding, MultiHeadAttention, Encoder & Decoder, 最后介绍如何把各个模块拼到一起。在这个过程中，我还会着重介绍一个论文里没有提及，但是代码实现时非常重要的一个细节——字符的处理。

说实话，用 PyTorch 实现 Transformer 没有什么有变数的地方，大家的代码写得都差不多，我也是参考着别人的教程写的。但是，Transformer 的代码实现中有很多坑。大部分人只会云淡风轻地介绍一下最终的代码成品，不会去讲他们 debug 耗费了多少时间，哪些地方容易出错。而我会着重讲一下代码中的一些细节，以及我碰到过的问题。

Positional Encoding

模型一开始是一个 Embedding 层加一个 Positional Encoding。Embedding 在 PyTorch 里已经有实现了，且不是文章的创新点，我们就直接来看 Positional Encoding 的写法。

求 Positional Encoding，其实就是求一个二元函数的许多函数值构成的矩阵。对于二元函数��(��,�)，我们要求出��∈[0,��−1],�∈[0,��−1]时所有的函数值，其中，��是该序列的长度，��是每一个词向量的长度。

理论上来说，每个句子的序列长度是不固定的，因此，每次都要用到不同的 Positional Encoding。但是，Positional Encoding 的值都是可以提前计算好的。我们可以设置一个很大的最大序列长度，预处理好这个超大的Positional Encoding矩阵。每次有句子输入进来，根据这个句子的序列长度，去预处理好的矩阵里取一小段出来即可。

这样，整个类的实现应该如下所示：

class PositionalEncoding(nn.Module):

    def __init__(self, d_model: int, max_seq_len: int):
        super().__init__()

        # Assume d_model is an even number for convenience
        assert d_model % 2 == 0

        i_seq = torch.linspace(1, max_seq_len, max_seq_len)
        j_seq = torch.linspace(2, d_model, d_model // 2)
        pos, two_i = torch.meshgrid(i_seq, j_seq)
        pe_2i = torch.sin(pos / 10000**(two_i / d_model))
        pe_2i_1 = torch.cos(pos / 10000**(two_i / d_model))
        pe = torch.stack((pe_2i, pe_2i_1), 2).reshape(1, max_seq_len, d_model)

        self.register_buffer('pe', pe, False)

    def forward(self, x: torch.Tensor):
        n, seq_len, d_model = x.shape
        pe: torch.Tensor = self.pe
        assert seq_len <= pe.shape[1]
        assert d_model == pe.shape[2]
        x *= d_model**0.5
        return x + pe[:, 0:seq_len, :]

代码中有不少需要讲解的部分。首先，先看一下预处理好的矩阵pe是怎么在__init__中算出来的。pe可以很直接地用两层循环算出来。由于这段预处理代码只会执行一次，相对于冗长的训练时间，哪怕生成pe的代码性能差一点也没事。然而，作为一个编程高手，我必须用高性能的并行编程把这段逻辑实现。

为了并行地求pe，我们要初始化一个二维网格，表示自变量��,�。生成网格可以用下面的代码实现。（由于�要分奇偶讨论，�的个数是��2)

i_seq = torch.linspace(1, max_seq_len, max_seq_len)
j_seq = torch.linspace(2, d_model, d_model // 2)
pos, two_i = torch.meshgrid(i_seq, j_seq)

torch.meshgrid用于生成网格。比如torch.meshgrid([0, 1], [0, 1])就可以生成[[(0, 0), (0, 1)], [(1, 0), (1, 1)]]这四个坐标构成的网格。不过，这个函数会把坐标的两个分量分别返回。比如：

i, j = torch.meshgrid([0, 1], [0, 1])
# i: [[0, 0], [1, 1]]
# j: [[0, 1], [0, 1]]

利用这个函数的返回结果，我们可以并行地分别算出奇偶位置的 PE 值。

pe_2i = torch.sin((pos / 10000)**(two_i / d_model))
pe_2i_1 = torch.cos((pos / 10000)**(two_i / d_model))

有了奇偶处的值，现在的问题是怎么把它们优雅地拼到同一个维度上。我这里先把它们堆成了形状为seq_len, d_model/2, 2的一个张量，再把最后一维展平，就得到了最后的pe矩阵。

pe = torch.stack((pe_2i, pe_2i_1), 2).reshape(1, max_seq_len, d_model)

最后，要注意一点。只用 self.pe = pe 记录这个量是不够好的。我们最好用 self.register_buffer('pe', pe, False) 把这个量登记成 torch.nn.Module 的一个存储区，同时也会完成self.pe = pe。这里涉及到 PyTorch 的一些知识了。

PyTorch 的 Module 会记录两类参数，一类是 parameter 可学习参数；另一类是 buffer 不可学习的参数。把变量登记成 buffer 的最大好处是，在使用 model.to(device) 把一个模型搬到另一个设备上时，所有 parameter 和 buffer 都会自动被搬过去。另外，buffer 和 parameter 一样，也可以被记录到 state_dict 中，并保存到文件里。register_buffer 的第三个参数决定了是否将变量加入 state_dict。由于 pe 可以直接计算，不需要记录，可以把这个参数设成 False。

预处理好 pe 后，用起来就很方便了。每次读取输入的序列长度，从中取一段出来即可。

另外，ransformer 给嵌入层乘了个系数��。为了方便起见，我把这个系数放到了 PositionalEncoding 类里面。

def forward(self, x: torch.Tensor):
    n, seq_len, d_model = x.shape
    pe: torch.Tensor = self.pe
    assert seq_len <= pe.shape[1]
    assert d_model == pe.shape[2]
    x *= d_model**0.5
    return x + pe[:, 0:seq_len, :]

Scaled Dot-Product Attention

下一步是多头注意力层。为了实现多头注意力，我们先要实现 Transformer 里经典的注意力计算。而在讲注意力计算之前，我还要补充一下 Transformer 中有关 mask 的一些知识。

Transformer 里的 mask

Transformer 最大的特点就是其并行性。给定翻译好的第1~n个词语，它默认会并行地预测第2~(n+1)个下一个词语。为了模拟串行输出的情况，第�个词语不应该看到第�+1个词语之后的信息。

输入信息	输出
(y1, --, --, --)	y2
(y1, y2, --, --)	y3
(y1, y2, y3, --)	y4
(y1, y2, y3, y4)	y5

为了实现这一功能，Transformer 使用了掩码。掩码取1表示这个地方的数是有效的，取0表示这个地方的数是无效的。Decoder 里的这种掩码应该是一个上三角全1矩阵。

掩码是在注意力计算中生效的。对于掩码取0的区域，其softmax前的��值取负无穷。这是因为，对于softmax

��=��1+��2+...��,

令��=−∞可以让它在 softmax 的分母里不产生任何贡献。

这是论文里提到的 mask，它用来模拟 decoder 的串行推理。而在代码实现中，还有其他地方会产生 mask。之前，我们在生成一个 batch 的数据时，给很多句子加入了。这个特殊字符也没有实际意义，不应该对计算产生任何贡献。因此，有的地方的 mask 也应该为0。之后讲 Transformer 模型类时，我会介绍所有的 mask 该怎么生成，这里我们仅关注注意力计算是怎么用到 mask 的。

注意力计算

补充完了背景知识，我们来看注意力计算的实现代码。由于注意力计算没有任何的状态，因此它应该写成一个函数，而不是一个类。我们可以轻松地用 PyTorch 代码翻译注意力计算的公式。（注意，我这里的 mask 表示哪些地方要填负无穷，而不是像之前讲的表示哪些地方有效）

��(�,�,�)=��(��)�

MY_INF = 1e12

def attention(q: torch.Tensor,
              k: torch.Tensor,
              v: torch.Tensor,
              mask: Optional[torch.Tensor] = None):
    '''
    # Note: The dtype of mask must be bool
    '''
    # q shape: [n, heads, q_len, d_k]
    # k v shape: [n, heads, k_len, d_v]
    assert q.shape[-1] == k.shape[-1]
    d_k = k.shape[-1]
    # tmp shape: [n, heads, q_len, k_len]
    tmp = torch.matmul(q, k.transpose(-2, -1)) / d_k**0.5
    if mask is not None:
        tmp.masked_fill_(mask, -MY_INF)
    tmp = F.softmax(tmp, -1)
    # tmp shape: [n, heads, q_len, d_v]
    tmp = torch.matmul(tmp, v)
    return tmp

这里有一个很坑的地方。引入了带来的 mask 后，会产生一个新的问题：可能一整行数据都是失效的，softmax 用到的所有 �� 可能都是负无穷。

��=�−∞�−∞+�−∞+...�−∞

这个数据是没有意义的。如果用torch.inf来表示无穷大，就会令exp(torch.inf)=0，最后 softmax 结果会出现 NaN。

按道理讲，若一行全是失效的数据，则它们的 softmax 结果应该全填零。所以，可以直接判断一下计算结果中有 NaN 的地方，把它们赋值成0。

但是，大多数 PyTorch Transformer 教程压根就没提这一点，而他们的代码又还是能够跑通。拿放大镜仔细对比了代码后，我发现，他们的无穷大用的不是 torch.inf，而是自己随手设的一个极大值。这样，exp(MY_INF)得到的不再是0，而是一个极小值。softmax 的结果就会等于分母的项数，而不是 NaN。讲道理，这些 softmax 结果是要填 0 的，填一个常数问题倒也不是很大。出于随大流的考虑，我也使用了这种懒散的实现方式，用了一个假的无穷大。

Multi-Head Attention

有了注意力计算，就可以实现多头注意力层了。多头注意力层是有学习参数的，它应该写成一个类。

class MultiHeadAttention(nn.Module):

    def __init__(self, heads: int, d_model: int, dropout: float = 0.1):
        super().__init__()

        assert d_model % heads == 0
        self.d_k = d_model // heads
        self.heads = heads
        self.d_model = d_model
        self.q = nn.Linear(d_model, d_model)
        self.k = nn.Linear(d_model, d_model)
        self.v = nn.Linear(d_model, d_model)
        self.out = nn.Linear(d_model, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self,
                q: torch.Tensor,
                k: torch.Tensor,
                v: torch.Tensor,
                mask: Optional[torch.Tensor] = None):
        # batch should be same
        assert q.shape[0] == k.shape[0]
        assert q.shape[0] == v.shape[0]
        # the sequence length of k and v should be aligned
        assert k.shape[1] == v.shape[1]

        n, q_len = q.shape[0:2]
        n, k_len = k.shape[0:2]
        q = self.q(q).reshape(n, q_len, self.heads, self.d_k).transpose(1, 2)
        k = self.k(k).reshape(n, k_len, self.heads, self.d_k).transpose(1, 2)
        v = self.v(v).reshape(n, k_len, self.heads, self.d_k).transpose(1, 2)

        attention_res = attention(q, k, v, mask)
        concat_res = attention_res.transpose(1, 2).reshape(
            n, q_len, self.d_model)
        concat_res = self.dropout(concat_res)

        output = self.out(concat_res)
        return output

这段代码一处很灵性的地方。在 Transformer 的论文中，多头注意力是先把每个词的表示拆成ℎ个头，再对每份做投影、注意力，最后拼接起来，再投影一次。其实，拆开与拼接操作是多余的。我们可以一开始就让所有头的数据经过同一个线性层，再在做注意力之前把头和序列数这两维转置一下。这两步操作和拆开来做投影、注意力是等价的。做完了注意力操作之后，再把两个维度转置回来，这和拼接操作是等价的。

def forward(self,
                q: torch.Tensor,
                k: torch.Tensor,
                v: torch.Tensor,
                mask: Optional[torch.Tensor] = None):
        # batch should be same
        assert q.shape[0] == k.shape[0]
        assert q.shape[0] == v.shape[0]
        # the sequence length of k and v should be aligned
        assert k.shape[1] == v.shape[1]

        n, q_len = q.shape[0:2]
        n, k_len = k.shape[0:2]
        q = self.q(q).reshape(n, q_len, self.heads, self.d_k).transpose(1, 2)
        k = self.k(k).reshape(n, k_len, self.heads, self.d_k).transpose(1, 2)
        v = self.v(v).reshape(n, k_len, self.heads, self.d_k).transpose(1, 2)

        attention_res = attention(q, k, v, mask)
        concat_res = attention_res.transpose(1, 2).reshape(
            n, q_len, self.d_model)
        concat_res = self.dropout(concat_res)

        output = self.out(concat_res)
        return output

前馈网络

前馈网络太简单了，两个线性层，没什么好说的。注意内部那个隐藏层的维度大小��会更大一点。

class FeedForward(nn.Module):

    def __init__(self, d_model: int, d_ff: int, dropout: float = 0.1):
        super().__init__()
        self.layer1 = nn.Linear(d_model, d_ff)
        self.dropout = nn.Dropout(dropout)
        self.layer2 = nn.Linear(d_ff, d_model)

    def forward(self, x):
        x = self.layer1(x)
        x = self.dropout(F.relu(x))
        x = self.layer2(x)
        return x

Encoder & Decoder

准备好一切组件后，就可以把模型一层一层搭起来了。先搭好每个 Encoder 层和 Decoder 层，再拼成 Encoder 和 Decoder。

Encoder 层和 Decoder 层的结构与论文中的描述一致，且每个子层后面都有一个 dropout，和上一层之间使用了残差连接。归一化的方法是 LayerNorm。顺带一提，不仅是这些层，前面很多子层的计算中都加入了 dropout。

再补充说明一下 mask。由于 encoder 和 decoder 的输入不同，它们的填充情况不同，产生的 mask 也不同。后文会展示这些 mask 的生成方法。

class EncoderLayer(nn.Module):

    def __init__(self,
                 heads: int,
                 d_model: int,
                 d_ff: int,
                 dropout: float = 0.1):
        super().__init__()
        self.self_attention = MultiHeadAttention(heads, d_model, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)

    def forward(self, x, src_mask: Optional[torch.Tensor] = None):
        x = self.norm1(x +
                       self.dropout1(self.self_attention(x, x, x, src_mask)))
        x = self.norm2(x + self.dropout2(self.ffn(x)))
        return x

class DecoderLayer(nn.Module):

    def __init__(self,
                 heads: int,
                 d_model: int,
                 d_ff: int,
                 dropout: float = 0.1):
        super().__init__()
        self.self_attention = MultiHeadAttention(heads, d_model, dropout)
        self.attention = MultiHeadAttention(heads, d_model, dropout)
        self.ffn = FeedForward(d_model, d_ff, dropout)
        self.norm1 = nn.LayerNorm(d_model)
        self.norm2 = nn.LayerNorm(d_model)
        self.norm3 = nn.LayerNorm(d_model)
        self.dropout1 = nn.Dropout(dropout)
        self.dropout2 = nn.Dropout(dropout)
        self.dropout3 = nn.Dropout(dropout)

    def forward(self,
                x,
                encoder_kv: torch.Tensor,
                dst_mask: Optional[torch.Tensor] = None,
                src_dst_mask: Optional[torch.Tensor] = None):
        x = self.norm1(x +
                       self.dropout1(self.self_attention(x, x, x, dst_mask)))
        x = self.norm2(x + self.dropout2(
            self.attention(x, encoder_kv, encoder_kv, src_dst_mask)))
        x = self.norm3(x + self.dropout2(self.ffn(x)))
        return x

Encoder 和 Decoder 就是再前面加了一个嵌入层、一个位置编码，再把多个层堆起来了而已，其他输入输出照搬。

class Encoder(nn.Module):

    def __init__(self,
                 vocab_size: int,
                 d_model: int,
                 d_ff: int,
                 n_layers: int,
                 heads: int,
                 dropout: float = 0.1,
                 max_seq_len: int = 120):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pe = PositionalEncoding(d_model, max_seq_len)
        self.layers = []
        for i in range(n_layers):
            self.layers.append(EncoderLayer(heads, d_model, d_ff, dropout))
        self.layers = nn.ModuleList(self.layers)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x, src_mask: Optional[torch.Tensor] = None):
        x = self.embedding(x)
        x = self.dropout(x)
        x = self.pe(x)
        x = self.dropout(x)
        for layer in self.layers:
            x = layer(x, src_mask)
        return x


class Decoder(nn.Module):

    def __init__(self,
                 vocab_size: int,
                 d_model: int,
                 d_ff: int,
                 n_layers: int,
                 heads: int,
                 dropout: float = 0.1,
                 max_seq_len: int = 120):
        super().__init__()
        self.embedding = nn.Embedding(vocab_size, d_model)
        self.pe = PositionalEncoding(d_model, max_seq_len)
        self.layers = []
        for i in range(n_layers):
            self.layers.append(DecoderLayer(heads, d_model, d_ff, dropout))
        self.layers = nn.Sequential(*self.layers)
        self.dropout = nn.Dropout(dropout)

    def forward(self,
                x,
                encoder_kv,
                dst_mask: Optional[torch.Tensor] = None,
                src_dst_mask: Optional[torch.Tensor] = None):
        x = self.embedding(x)
        x = self.dropout(x)
        x = self.pe(x)
        x = self.dropout(x)
        for layer in self.layers:
            x = layer(x, encoder_kv, dst_mask, src_dst_mask)
        return x

Transformer 类

终于，激动人心的时候到来了。我们要把各个子模块组成变形金刚（Transformer）了。先过一遍所有的代码。

class Transformer(nn.Module):

    def __init__(self,
                 src_vocab_size: int,
                 dst_vocab_size: int,
                 d_model: int,
                 d_ff: int,
                 n_layers: int,
                 heads: int,
                 dropout: float = 0.1,
                 max_seq_len: int = 200):
        super().__init__()
        self.encoder = Encoder(src_vocab_size, d_model, d_ff, n_layers, heads,
                               dropout, max_seq_len)
        self.decoder = Decoder(dst_vocab_size, d_model, d_ff, n_layers, heads,
                               dropout, max_seq_len)
        self.heads = heads

    def init_weights(self):
        for p in self.parameters():
            if p.dim() > 1:
                nn.init.xavier_uniform_(p)

    def generate_mask(self,
                      q_pad: torch.Tensor,
                      k_pad: torch.Tensor,
                      with_left_mask: bool = False):
        # q_pad shape: [n, q_len]
        # k_pad shape: [n, k_len]
        # q_pad k_pad dtype: bool
        assert q_pad.device == k_pad.device
        n, q_len = q_pad.shape
        n, k_len = k_pad.shape

        mask_shape = (n, self.heads, q_len, k_len)
        if with_left_mask:
            mask = 1 - torch.tril(torch.ones(mask_shape))
        else:
            mask = torch.zeros(mask_shape)
        mask = mask.to(q_pad.device)
        for i in range(n):
            mask[i, :, q_pad[i], :] = 1
            mask[i, :, :, k_pad[i]] = 1
        mask = mask.to(torch.bool)
        return mask

    def forward(self,
                x,
                y,
                src_pad_mask: Optional[torch.Tensor] = None,
                dst_pad_mask: Optional[torch.Tensor] = None):
        src_mask = self.generate_mask(src_pad_mask, src_pad_mask, False)
        dst_mask = self.generate_mask(dst_pad_mask, dst_pad_mask, True)
        src_dst_mask = self.generate_mask(dst_pad_mask, src_pad_mask, False)
        encoder_kv = self.encoder(x, src_mask)
        res = self.decoder(y, encoder_kv, dst_mask, src_dst_mask)
        embedding_reverse = self.decoder.embedding.weight.transpose(0, 1)
        res = torch.matmul(res, embedding_reverse)
        return res

我们一点一点来看。先看初始化函数。初始化函数的输入其实就是 Transformer 模型的超参数。总结一下，Transformer 应该有这些超参数：

d_model 模型中大多数词向量表示的维度大小
d_ff 前馈网络隐藏层维度大小
n_layers 堆叠的 Encoder & Decoder 层数
head 多头注意力的头数
dropout Dropout 的几率

另外，为了构件嵌入层，要知道源语言、目标语言的词典大小。为了预处理位置编码，需要提前知道一个最大序列长度。

照着子模块的初始化参数表，把参数归纳到__init__的参数表里即可。

def __init__(self,
            src_vocab_size: int,
            dst_vocab_size: int,
            d_model: int,
            d_ff: int,
            n_layers: int,
            heads: int,
            dropout: float = 0.1,
            max_seq_len: int = 200):
    super().__init__()
    self.encoder = Encoder(src_vocab_size, d_model, d_ff, n_layers, heads,
                            dropout, max_seq_len)
    self.decoder = Decoder(dst_vocab_size, d_model, d_ff, n_layers, heads,
                            dropout, max_seq_len)
    self.heads = heads

再看一下 forward 函数。forward先预处理好了所有的 mask，再逐步执行 Transformer 的计算：先是通过 Encoder 获得源语言的中间表示，再把它和目标语言的输入一起传入 Decoder，最后经过线性层输出结果。由于 PyTorch 的交叉熵损失函数自带了 softmax 操作，这里不需要多此一举。

当然，你去对着上面这段描述看下面的代码，会感到疑惑：输出前哪有线性层？embedding_reverse 又是什么东西？Transformer 论文提到，softmax 前的那个线性层可以和嵌入层共享权重，也就是说它们分别完成了词序号到词嵌入的正反映射。因此，我尊重原论文，用嵌入层的权重的转置完成了输出前的线性运算。不知道为什么网上大多数的实现都重新定义了一个线性层，可能是二者的区别不大吧。

def forward(self,
            x,
            y,
            src_pad_mask: Optional[torch.Tensor] = None,
            dst_pad_mask: Optional[torch.Tensor] = None):
    src_mask = self.generate_mask(src_pad_mask, src_pad_mask, False)
    dst_mask = self.generate_mask(dst_pad_mask, dst_pad_mask, True)
    src_dst_mask = self.generate_mask(dst_pad_mask, src_pad_mask, False)
    encoder_kv = self.encoder(x, src_mask)
    res = self.decoder(y, encoder_kv, dst_mask, src_dst_mask)
    embedding_reverse = self.decoder.embedding.weight.transpose(0, 1)
    res = torch.matmul(res, embedding_reverse)
    return res

等了很久，现在可以来仔细看一看 mask 的生成方法了。回忆一下，表示该字符是否有效的 mask 有两个来源。第一个是论文里提到的，用于模拟串行推理的 mask；另一个是填充操作的空白字符引入的 mask。generate_mask 用于生成这些 mask。

generate_mask 的输入有 query 句子和 key 句子的 pad mask q_pad, k_pad，它们的形状为[n, seq_len]。若某处为 True，则表示这个地方的字符是。对于自注意力，query 和 key 都是一样的；而在 Decoder 的第二个多头注意力层中，query 来自目标语言，key 来自源语言。with_left_mask 表示是不是要加入 Decoder 里面的那种 mask。

def generate_mask(self,
                q_pad: torch.Tensor,
                k_pad: torch.Tensor,
                with_left_mask: bool = False):

一开始，先取好维度信息，定好张量的形状。在注意力操作中，softmax 前的那个量的形状是 [n, heads, q_len, k_len]，表示每一批每一个头的每一个query对每个key之间的相似度。mask 的形状应和它一样。

再新建一个表示最终 mask 的张量。如果不用 Decoder 的那种 mask，就生成一个全零的张量；否则，生成一个上三角为0，其余地方为1的张量。注意，在我的代码中，mask 为 True 或1就表示这个地方需要填负无穷。

最后，把有的地方也标记一下。从形状[n, heads, q_len, k_len]可以知道，q_pad 表示哪些行是无效的，k_pad 表示哪些列是无效的。把对应的行或列赋值成1即可。

def generate_mask(self,
                    q_pad: torch.Tensor,
                    k_pad: torch.Tensor,
                    with_left_mask: bool = False):
    # q_pad shape: [n, q_len]
    # k_pad shape: [n, k_len]
    # q_pad k_pad dtype: bool
    assert q_pad.device == k_pad.device
    n, q_len = q_pad.shape
    n, k_len = k_pad.shape

    mask_shape = (n, self.heads, q_len, k_len)
    if with_left_mask:
        mask = 1 - torch.tril(torch.ones(mask_shape))
    else:
        mask = torch.zeros(mask_shape)
    mask = mask.to(q_pad.device)
    for i in range(n):
        mask[i, :, q_pad[i], :] = 1
        mask[i, :, :, k_pad[i]] = 1
    mask = mask.to(torch.bool)
    return mask

在 Transformer 中，有三类多头注意力层，它们的 mask 也不同。Encoder 的多头注意力层的 query 和 key 都来自源语言；Decoder 的第一个多头注意力层的 query 和 key 都来自目标语言；Decoder 的第二个多头注意力层的 query 来自目标语言， key 来自源语言。另外，Decoder 的第一个多头注意力层要加串行推理的那个 mask。

def forward(self,
            x,
            y,
            src_pad_mask: Optional[torch.Tensor] = None,
            dst_pad_mask: Optional[torch.Tensor] = None):
    src_mask = self.generate_mask(src_pad_mask, src_pad_mask, False)
    dst_mask = self.generate_mask(dst_pad_mask, dst_pad_mask, True)
    src_dst_mask = self.generate_mask(dst_pad_mask, src_pad_mask, False)

到此，Transfomer 模型总算编写完成了。噢，别忘了，模型参数的初始化方法也可以写进类里面。

def init_weights(self):
    for p in self.parameters():
        if p.dim() > 1:
            nn.init.xavier_uniform_(p)

训练

准备好了模型、数据集后，剩下的工作非常惬意，只要随便调用一下就行了。训练的代码如下：

# Config
batch_size = 64
lr = 0.0001
d_model = 256
d_ff = 1024
n_layers = 6
heads = 8


def main():
    en_vocab, zh_vocab = load_vocab()

    en_train, zh_train, en_valid, zh_valid = load_sentences()
    dataloader_train = get_dataloader(en_train, zh_train, batch_size)

    device = 'cuda:0'

    print_interval = 1000

    model = Transformer(len(en_vocab), len(zh_vocab), d_model, d_ff, n_layers,
                        heads)
    model.to(device)

    model.init_weights()

    optimizer = torch.optim.Adam(model.parameters(), lr)
    citerion = nn.CrossEntropyLoss(ignore_index=PAD_ID)
    tic = time.time()
    cnter = 0
    dataset_len = len(dataloader_train.dataset)
    print('Dataset size:', dataset_len)
    for epoch in range(10):
        loss_sum = 0

        for x, y in dataloader_train:
            x, y = x.to(device), y.to(device)
            x_mask = x == PAD_ID
            y_mask = y == PAD_ID
            y_input = y[:, :-1]
            y_label = y[:, 1:]
            y_mask = y_mask[:, :-1]
            y_hat = model(x, y_input, x_mask, y_mask)
            n, seq_len = y_label.shape
            y_hat = torch.reshape(y_hat, (n * seq_len, -1))
            y_label = torch.reshape(y_label, (n * seq_len, ))
            loss = citerion(y_hat, y_label)

            optimizer.zero_grad()
            loss.backward()
            torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
            optimizer.step()

            loss_sum += loss.item()

            toc = time.time()
            interval = toc - tic
            minutes = int(interval // 60)
            seconds = int(interval % 60)
            if cnter % print_interval == 0:
                print(f'{cnter:08d} {minutes:02d}:{seconds:02d}'
                      f' loss: {loss.item()}')
            cnter += 1

        print(f'Epoch {epoch}. loss: {loss_sum / dataset_len}')

    torch.save(model.state_dict(), 'dldemos/Transformer/model.pth')
    print('Done.')

准备好超参数，并统计好词典的大小，就可以初始化模型结构及权重了。

# Config
batch_size = 64
lr = 0.0001
d_model = 256
d_ff = 1024
n_layers = 6
heads = 8

model = Transformer(len(en_vocab), len(zh_vocab), d_model, d_ff, n_layers,
                    heads)
model.to(device)

model.init_weights()

优化器用 Adam，非常普通。损失函数是交叉熵函数。这里需要注意一下，不仅模型计算中需要忽略掉，算损失函数时也需要扣除它们的贡献。CrossEntropyLoss的ignore_index表示标签中的哪一类需要被忽略。这里我们传入字符在词典中的序号即可。

optimizer = torch.optim.Adam(model.parameters(), lr)
citerion = nn.CrossEntropyLoss(ignore_index=PAD_ID)

模型训练的部分也比较简单。先预处理出输入里的 mask 和模型的输入输出。如前所述，decoder的输入是第一个到倒数第二个字符y[:, :-1]，期待的输出是第二到最后一个字符y[:, 1:]。

再来就是模型前向传播与误差计算了。注意，由于数据带了序列长度那一维，在传入交叉熵之前要把那一维和 batch 维合并。

最后梯度下降的操作也很常规。注意截断一下梯度，防止梯度爆炸即可。

for x, y in dataloader_train:
    x, y = x.to(device), y.to(device)
    x_mask = x == PAD_ID
    y_mask = y == PAD_ID
    y_input = y[:, :-1]
    y_label = y[:, 1:]
    y_mask = y_mask[:, :-1]
    y_hat = model(x, y_input, x_mask, y_mask)
    n, seq_len = y_label.shape
    y_hat = torch.reshape(y_hat, (n * seq_len, -1))
    y_label = torch.reshape(y_label, (n * seq_len, ))
    loss = citerion(y_hat, y_label)

    optimizer.zero_grad()
    loss.backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
    optimizer.step()

由于 Transformer 训练得较慢，得过很久才能看到训练结果。在调试时，我们只要观察到 loss 在不断下降，就可以大胆宣告胜利了。

实验

待更新

总结

只要读懂了 Transfomer 的论文，用 PyTorch 实现一遍 Transformer 是很轻松的。但是，代码实现中有非常多论文不会提及的细节，预处理数据集也算是一个很烦人的工作。在这篇文章里，我完整地介绍了一个英中翻译 Transformer 的 PyTorch 实现，相信读者能够跟随这篇文章实现自己的 Transformer，并在代码实现的过程中加深对论文的理解。

再稍微总结一下代码实现中的一些值得注意的地方。代码中最大的难点是 mask 的实现。mask 的处理稍有闪失，就可能会让计算结果中遍布 NaN。一定要想清楚各个模块的 mask 是从哪来的，它们在注意力计算里是怎么被用上的。

另外，有两处地方的实现比较灵活。一处是位置编码的实现，一处是多头注意力中怎么描述“多头”。其他模块的实现都大差不差，千篇一律。

至此，我的《吴恩达深度学习专项代码实战》总算是完结了。过几天来写总结。

参考资料

https://github.com/hyunwoongko/transformer

https://towardsdatascience.com/

原文：https://zhuanlan.zhihu.com/p/581334630

PyTorch Transformer 英中翻译超详细教程

任务与数据集准备

数据集预处理

数据清洗

词语转序号

生成 Dataloader

Transformer 模型

Positional Encoding

Scaled Dot-Product Attention

Transformer 里的 mask

注意力计算

Multi-Head Attention

前馈网络

Encoder & Decoder

Transformer 类

训练

实验

总结

参考资料

语音

热门讨论

热门文章

相关电子书

PyTorch Transformer 英中翻译超详细教程

任务与数据集准备

数据集预处理

数据清洗

词语转序号

生成 Dataloader

Transformer 模型

Positional Encoding

Scaled Dot-Product Attention

Transformer 里的 mask

注意力计算

Multi-Head Attention

前馈网络

Encoder & Decoder

Transformer 类

训练

实验

总结

参考资料

语音

热门讨论

热门文章

相关文章

相关电子书

相关镜像