基于Transformer实现机器翻译

前言

  在上一篇博客中，我们已经详细分析了Encoder-Decoder架构和Attention机制，并基于二者实现了机器翻译的实例，本文将在此基础上，进一步使用Transformer进行机器翻译。

一、Transformer详解

  Transformer是一种基于注意力机制的神经网络架构，最初由Vaswani等人在2017年的论文《Attention is All You Need》中提出。它广泛应用于NLP任务，如机器翻译、文本生成和文本分类。Transformer通过注意力机制，摒弃了传统的RNN和CNN，实现了更高效的计算和更好的性能。

Transformer结构图

  上图是论文中 Transformer 的内部结构图，左侧为 Encoder block，右侧为 Decoder block。可以看到 Encoder block 包含一个 Multi-Head Attention，而 Decoder block 包含两个 Multi-Head Attention (其中有一个用到 Masked)。Multi-Head Attention 上方还包括一个 Add & Norm 层，Add 表示残差连接 (Residual Connection) 用于防止网络退化，Norm 表示 Layer Normalization，用于对每一层的激活值进行归一化。下面，我们将从输入到输出，逐步分析Transformer的结构和计算过程。

1.1 Transformer的Inputs

1.1.1 输入嵌入（Input Embedding）

  输入序列首先通过嵌入层（Embedding Layer）转化为固定维度的向量表示。嵌入层的作用是将每个词转换为高维空间中的向量表示，这样可以捕捉到词与词之间的语义关系。

1.1.2 位置编码（Positional Encoding）

  由于Transformer不具备RNN和CNN那样的顺序处理能力，位置编码用于引入序列位置信息，帮助模型捕捉顺序信息。位置编码通过正弦和余弦函数生成，公式如下：

  其中，表示单词在句子中的绝对位置，表示维度索引，表示嵌入向量的维度。

  最终，输入嵌入和位置编码相加，得到带有位置信息的输入向量。(注意：为了减少维度、加速训练，应该将positional encoding与词向量相加，而不是拼接)

  Transformer 的 Decoder的输入与Encoder的输出处理方法步骤相同，只是一个接受source数据，一个接受target数据，例如在翻译过程中：Encoder接受英文"Tom chase Jerry"，则Decoder应该接受中文"汤姆追逐杰瑞"。仅在有target数据时（也就是有监督训练时）Decoder 才会接受Outputs Embedding，进行预测时则不会接收。

1.2 Transformer的Encoder

  编码器由N个相同的层组成，每层包括多头自注意力机制（Multi-Head Self-Attention Mechanism）和前馈神经网络。

Encoder结构

1.2.1 自注意力机制（Self-Attention Mechanism）

  随着模型处理输入序列的每个单词，自注意力会关注整个输入序列的所有单词，帮助模型对本单词更好地进行编码。在处理过程中，自注意力机制会将对所有相关单词的理解融入到我们正在处理的单词中。具体的功能如下：

（1）序列建模：自注意力可以用于序列数据（例如文本、时间序列、音频等）的建模。它可以捕捉序列中不同位置的依赖关系，从而更好地理解上下文。这对于机器翻译、文本生成、情感分析等任务非常有用。

（2）并行计算：自注意力可以并行计算，这意味着可以有效地在现代硬件上进行加速。相比于RNN和CNN等序列模型，它更容易在GPU和TPU等硬件上进行高效的训练和推理。（因为在自注意力中可以并行的计算得分）

（3）长距离依赖捕捉：传统的循环神经网络（RNN）在处理长序列时可能面临梯度消失或梯度爆炸的问题。自注意力可以更好地处理长距离依赖关系，因为它不需要按顺序处理输入序列。

Self-Attention结构

  自注意力机制通过计算输入序列中每个位置与其他位置的关系，生成权重并进行加权求和。具体过程如下：

  1. 从每个编码器的输入向量中生成三个向量： 将输入向量通过三个不同的线性变换，生成查询向量（Query，）、键向量（Key，）和值向量（Value，）。

  2. 计算得分： 计算查询向量和键向量的点积，以衡量输入序列中各个位置之间的相关性。公式如下：

  3. 缩放得分： 将得分除以键向量维度的平方根（通常为8，这是论文中使用的键向量维度64的平方根）。这样做是为了防止内积过大，使梯度更稳定。公式如下：

      其中， 是键向量的维度。

  4. 通过 softmax 传递结果： 对缩放后的得分进行 softmax 变换，以获得注意力权重。这些权重表示输入序列中每个位置对当前处理位置的重要程度。公式如下：

  5. 加权值向量： 将每个值向量乘以相应的 softmax 权重。这一步是为了关注语义上相关的单词，并弱化不相关的单词。公式如下：

  6. 对加权值向量求和： 对加权后的值向量进行加权求和，得到最终的自注意力输出。

      综上所述，Self-Attention的计算公式为：

自注意力-计算示例

  通过上述步骤，即可实现自注意力的计算

1.2.2 多头注意力机制（Multi-Head Attention Mechanism）

  Multi-Head Attention 在 self-attention 的基础上，对于输入的 embedding 矩阵通过多个不同的头并行计算注意力，目的是捕获不同子空间的特征。self-attention 只使用了一组 , , 来进行变换得到 Query、Keys、Values。而 Multi-Head Attention 使用多组 , , 得到多组 Query、Keys、Values，然后每组分别计算得到一个 Z 矩阵，最后将得到的多个 Z 矩阵进行拼接。Transformer 用了 8 组.

1.2.3 前馈神经网络（Feed-Forward Neural Network）

  前馈神经网络结构比较简单，包括两个线性变换和一个ReLU激活函数：

1.2.4 Add & Norm 详解

  另外，每个子层（注意力机制和前馈网络）后面都有Add & Norm层。其主要包括两部分，即：残差连接（Residual Connection）和层归一化（Layer Normalization）。

（1）残差连接（Residual Connection）

  残差连接的目的是为了缓解深层网络中可能出现的梯度消失问题。通过将输入直接传递到输出，模型可以更容易地学习恒等映射，这样即使某些层没有学习到有用的信息，也不会对整个网络造成严重影响。

（2）层归一化（Layer Normalization）

  层归一化用于稳定训练过程，加速收敛。它通过对每个样本的特征进行归一化，使得输出具有零均值和单位方差，从而消除不同特征之间的尺度差异。

  具体来说，对于输入 ，层归一化计算如下：

  其中， 和 分别是 的均值和方差， 是一个小常数，用于防止除零错误， 和 是可学习的参数，用于重新缩放和偏移归一化后的结果。

1.3 Transformer的Decoder

  解码器同样由N个相同的层组成，每层包括掩码多头自注意力机制（Masked Multi-Head Self-Attention Mechanism）、编码器-解码器多头注意力机制（Encoder-Decoder Attention Mechanism）和前馈神经网络。

Encoder-Decoder

  在学习了Encoder之后，Decoder的结构就比较清晰易懂了。由于前馈神经网络和 Add & Norm 层与Encoder的部分相同，下文不再重复，重点介绍Decoder中的注意力机制。

1.3.1 解码器的输入

  解码器的输入分为两类：一种是训练时的输入，一种是预测时的输入。在训练时，解码器输入为实际的目标序列；在预测时，解码器输入为已经生成的序列。

1.3.2 掩码多头自注意力机制（Masked Multi-Head Self-Attention Mechanism）

  Masked Multi-Head Attention 的计算原理与 Encoder 的 Multi-Head Attention 相同，只是多加了一个掩码（mask）。掩码用于屏蔽某些值，使其在参数更新时不产生效果。Transformer 模型中涉及两种掩码：padding mask 和 sequence mask。那么为什么需要添加这两种掩码呢？

（1）Padding Mask

  因为每个批次输入序列的长度不一样，我们需要对输入序列进行对齐。具体来说，就是在较短的序列后面填充 0。但是，如果输入的序列太长，则截取左边的内容，多余部分直接舍弃。因为这些填充的位置没有实际意义，所以注意力机制不应该关注这些位置。为此，我们需要将这些位置的值加上一个非常大的负数（负无穷），这样经过 softmax 变换，这些位置的概率就会接近 0。

（2）Sequence Mask

  sequence mask 是为了使解码器（decoder）不能看到未来的信息。对于一个序列，在时间步 t 时刻，解码输出应该只能依赖于 t 时刻之前的输出，而不能依赖 t 之后的输出。因此，我们需要隐藏 t 之后的信息。这在训练时尤其重要，因为训练时我们将完整的目标数据输入到解码器中，而预测时则不需要，在预测时我们只能依赖前一时刻的输出。具体做法是构建一个上三角矩阵，上三角的值全为 0，把这个矩阵作用在每一个序列上，即可实现。

  注意：在 Encoder 中的 Multi-Head Attention 也是需要进行掩码处理的，只不过 Encoder 中只需要 padding mask，而解码器中需要同时使用 padding mask 和 sequence mask。除了这点不同，其他部分与 Encoder 中的 Multi-Head Attention 相同。

1.3.3 编码器-解码器多头注意力机制（Encoder-Decoder Attention Mechanism）

  使用编码器的输出作为键和值，解码器的输出作为查询，计算注意力，帮助解码器在生成序列时参考编码器的上下文信息。

  在解码器中的第二个 Multi-Head Attention 与 Encoder 中的有些不同。Encoder 中的 Multi-Head Attention 是基于 Self-Attention 的，而解码器中的第二个 Multi-Head Attention 仅基于 Attention。它的查询向量（Query）来自于 Masked Multi-Head Attention 的输出，键向量（Key）和值向量（Value）来自于 Encoder 中最后一层的输出。

  为什么解码器中需要两个 Multi-Head Attention？我的理解是，第一个 Masked Multi-Head Attention 用于捕捉之前已经预测的输出信息，相当于记录当前时刻输入之间的信息。第二个 Multi-Head Attention 则是通过当前输入的信息得到下一时刻的信息，也就是输出的信息，用于表示当前的输入与 Encoder 提取的特征向量之间的关系，从而预测输出。

  经过第二个 Multi-Head Attention 之后，处理流程与 Encoder 相同。然后，输出进入下一个解码器层，经过多层解码器的处理后，最终到达输出层。

1.4 Transformer的Output

  解码器的输出通过线性变换映射到词汇表大小的向量，并通过Softmax变换得到输出的概率分布：

  最终通过词典，输出概率最大的对应的单词作为我们的预测输出。

1.5 模型特点总结

（1）Transformer摒弃了传统的RNN和CNN结构，通过注意力机制实现了全局信息的高效捕捉和传递。

（2）编码器和解码器中的各个层可以并行计算，提高了训练和推理速度。

（3）多头注意力机制能够同时关注序列中的不同部分，有效捕捉长距离依赖关系。

（4）残差连接和层归一化确保了模型训练的稳定性和快速收敛。

二、Transformer机器翻译实例

  日文和中文之间的翻译，使用 PyTorch、Torchtext 和 SentencePiece 实现   作者在AutoDL算力云（https://www.autodl.com） 租用设备上进行训练，展示了GPU训练过程和结果图

实验环境

实验环境

2.1 导入库并下载数据集

  在本小节中，我们将使用从 JParaCrawl[http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl] 下载的日英并行数据集。该数据集被称为 "NTT 创建的最大的公开可用英日并行语料库"。它主要是通过自动对齐平行句子的网络爬虫创建的。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
import math
import torchtext
import torch
import torch.nn as nn
from torch import Tensor
from torch.nn.utils.rnn import pad_sequence
from torch.utils.data import DataLoader
from collections import Counter
from torchtext.vocab import Vocab
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
import io
import time
import pandas as pd
import numpy as np
import pickle
import tqdm
import sentencepiece as spm
torch.manual_seed(0)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# print(torch.cuda.get_device_name(0)) 
device

# 读取数据，使用Tab字符('\t')作为分隔符
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainen = df[2].values.tolist()#[:10000] # 英文训练数据
trainja = df[3].values.tolist()#[:10000] # 日文训练数据
# trainen.pop(5972)
# trainja.pop(5972)

  在导入所有日语及其对应英语后，我们删除了数据集中的最后一个数据，因为它有一个缺失值。总的来说，trainen 和 trainja 中的句子数量为 5,973,071 个。不过，出于学习目的，通常建议在一次性使用所有数据之前，先对数据进行采样，以确保一切正常节省时间。

  下面是数据集中的一个句子示例。 
1
2
3
# 打印训练数据列表中的第500个元素
print(trainen[500])
print(trainja[500])

train[500]

train[500]

  我们也可以使用不同的并行数据集来跟进本文，只需确保我们能将数据处理成如上图所示的两个字符串列表，其中包含日语和英语句子。

2.2 准备分词器

  与英语或其他字母语言不同，日语句子不包含用来分隔单词的空格。我们可以使用由JParaCrawl提供的分词器，这些分词器是使用SentencePiece为日语和英语创建的。您可以访问JParaCrawl网站下载这些分词器。

1
2
3
# 加载英文、日文的SentencePiece分词模型
en_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.ja.nopretok.model')

  加载分词器后，您可以通过执行以下代码来测试它们。

1
2
# 使用英文的SentencePiece分词模型对给定句子进行编码
en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.")

英语编码

英语编码 
1
2
# 使用日文的SentencePiece分词模型对给定句子进行编码
ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。")

日语编码

日语编码

2.3 构建TorchText的Vocab对象并将句子转换为Torch张量

  使用分词器和原始句子，然后构建从TorchText导入的Vocab对象。这个过程可能需要几秒钟或几分钟，具体取决于数据集的大小和计算能力。不同的分词器也会影响构建词汇表所需的时间，我尝试了几种其他的日语分词器，但SentencePiece对本例而言效果更好且速度足够快。

1
2
3
4
5
6
7
8
9
10
11
def build_vocab(sentences, tokenizer):
    counter = Counter()  # 初始化计数器
    for sentence in sentences:
        # 使用分词器对句子进行编码，并更新计数器
        counter.update(tokenizer.encode(sentence, out_type=str))
    # 构建词汇表，并添加特殊标记
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])

# 使用日文、英文训练数据和分词器构建日文词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)
en_vocab = build_vocab(trainen, en_tokenizer)

  拥有词汇对象后，我们可以使用词汇和分词器对象来为我们的训练数据构建张量。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
def data_process(ja, en):
    data = []
    # 遍历日文和英文的训练数据
    for (raw_ja, raw_en) in zip(ja, en):
        # 使用日文、英文词汇表将日文句子编码为张量
        ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                                  dtype=torch.long)
        en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)],
                                  dtype=torch.long)
        # 将编码后的张量对添加到数据列表中
        data.append((ja_tensor_, en_tensor_))
    return data

# 处理训练数据，生成编码后的张量对
train_data = data_process(trainja, trainen)

2.4 创建在训练期间迭代的DataLoader对象

  在这里，我将BATCH_SIZE设置为16，以防止“cuda out of memory”错误，但这取决于机器内存容量、数据大小等各种因素，因此可以根据需要随意更改批量大小（注意：PyTorch教程使用Multi30k德英数据集时，将批量大小设置为128）。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
BATCH_SIZE = 8  # 设置批量大小
PAD_IDX = ja_vocab['<pad>']  # 获取日文词汇表中'<pad>'的索引
BOS_IDX = ja_vocab['<bos>']  # 获取'<bos>'的索引
EOS_IDX = ja_vocab['<eos>']  # 获取'<eos>'的索引

def generate_batch(data_batch):
    ja_batch, en_batch = []  # 初始化日文和英文批次列表
    for (ja_item, en_item) in data_batch:
        # 为日文、英文序列添加BOS和EOS标记
        ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
        en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
    # 填充日文、英文序列，使其长度一致
    ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
    en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)
    return ja_batch, en_batch  # 返回填充后的日文和英文批次

# 创建数据加载器，按批次加载训练数据，打乱顺序，并使用generate_batch函数进行批次生成
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

2.5 Sequence-to-sequence Transformer

  接下来的几段代码和文字解释（用斜体书写）取自原始的PyTorch教程 [https://pytorch.org/tutorials/beginner/translation_transformer.html]. 除了BATCH_SIZE和将de_vocab改为ja_vocab外，我没有做任何更改。

  Transformer是“Attention is all you need”论文中引入的用于解决机器翻译任务的Seq2Seq模型。Transformer模型由编码器和解码器块组成，每个块包含固定数量的层。

  编码器通过多头注意力和前馈网络层系列来处理输入序列。编码器的输出称为记忆（memory），与目标张量一起输入到解码器。编码器和解码器通过教师强制技术（teacher forcing technique）进行端到端训练。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)

class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, src_vocab_size: int, tgt_vocab_size: int,
                 dim_feedforward: int = 512, dropout: float = 0.1):
        super(Seq2SeqTransformer, self).__init__()
        # 初始化Transformer编码器层
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 堆叠多个编码器层，构建Transformer编码器
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        # 初始化Transformer解码器层
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 堆叠多个解码器层，构建Transformer解码器
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)
        
        # 定义输出生成器，将解码器的输出映射到目标词汇表大小
        self.generator = nn.Linear(emb_size, tgt_vocab_size)
        # 定义源语言的词嵌入层
        self.src_tok_emb = TokenEmbedding(src_vocab_size, emb_size)
        # 定义目标语言的词嵌入层
        self.tgt_tok_emb = TokenEmbedding(tgt_vocab_size, emb_size)
        # 定义位置编码层
        self.positional_encoding = PositionalEncoding(emb_size, dropout=dropout)

    def forward(self, src: Tensor, trg: Tensor, src_mask: Tensor,
                tgt_mask: Tensor, src_padding_mask: Tensor,
                tgt_padding_mask: Tensor, memory_key_padding_mask: Tensor):
        # 对源序列进行词嵌入和位置编码
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        # 对目标序列进行词嵌入和位置编码
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        # 通过Transformer编码器
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        # 通过Transformer解码器
        outs = self.transformer_decoder(tgt_emb, memory, tgt_mask, None,
                                        tgt_padding_mask, memory_key_padding_mask)
        # 生成最终的输出
        return self.generator(outs)

    def encode(self, src: Tensor, src_mask: Tensor):
        # 编码器前向传播
        return self.transformer_encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        # 解码器前向传播
        return self.transformer_decoder(self.positional_encoding(
                          self.tgt_tok_emb(tgt)), memory,
                          tgt_mask)

  文本标记通过使用标记嵌入来表示。位置编码被添加到标记嵌入中，以引入单词顺序的概念。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        super(PositionalEncoding, self).__init__()
        # 计算位置编码中的指数项
        den = torch.exp(-torch.arange(0, emb_size, 2) * math.log(10000) / emb_size)
        # 生成位置索引
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        # 初始化位置编码矩阵
        pos_embedding = torch.zeros((maxlen, emb_size))
        # 为位置编码矩阵的偶数列赋值
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        # 为奇数列赋值
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        # 在最后添加一个维度，以便与token嵌入相加
        pos_embedding = pos_embedding.unsqueeze(-2)

        # 定义Dropout层
        self.dropout = nn.Dropout(dropout)
        # 注册位置编码矩阵为缓冲区，不作为模型参数进行训练
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        # 将位置编码添加到token嵌入中，添加Dropout
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0), :])

class TokenEmbedding(nn.Module):
    def __init__(self, vocab_size: int, emb_size):
        super(TokenEmbedding, self).__init__()
        # 定义词嵌入层
        self.embedding = nn.Embedding(vocab_size, emb_size)
        self.emb_size = emb_size
    def forward(self, tokens: Tensor):
        # 对token进行嵌入，并乘以嵌入维度的平方根
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

  我们创建了一个后续词遮罩（word mask），以防止目标词关注其后续词。我们还创建了用于遮盖源和目标填充标记的mask。 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def generate_square_subsequent_mask(sz):
    # 生成一个大小为(sz, sz)的上三角矩阵
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将上三角矩阵转换为float类型，并将非上三角部分填充为负无穷
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

def create_mask(src, tgt):
    # 获取源序列和目标序列的长度
    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]

    # 为目标序列生成一个大小为(tgt_seq_len, tgt_seq_len)的方形掩码
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)
    # 为源序列生成一个全零矩阵的掩码
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)

    # 生成源序列和目标序列的填充掩码
    src_padding_mask = (src == PAD_IDX).transpose(0, 1)
    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)
    
    return src_mask, tgt_mask, src_padding_mask, tgt_padding_mask

  定义模型参数并实例化模型： 

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
SRC_VOCAB_SIZE = len(ja_vocab)  # 源语言词汇表大小
TGT_VOCAB_SIZE = len(en_vocab)  # 目标语言词汇表大小
EMB_SIZE = 512  # 词嵌入维度
NHEAD = 8  # 多头注意力机制的头数
FFN_HID_DIM = 512  # 前馈神经网络的隐藏层维度
BATCH_SIZE = 16  # 批量大小
NUM_ENCODER_LAYERS = 3  # 编码器层数
NUM_DECODER_LAYERS = 3  # 解码器层数
NUM_EPOCHS = 16  # 训练轮数

# 初始化Seq2SeqTransformer模型
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS,
                                 EMB_SIZE, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,
                                 FFN_HID_DIM)

# 用Xavier初始化方法初始化模型参数
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)

# 将模型移动到设备上（GPU）
transformer = transformer.to(device)

# 定义损失函数，忽略填充索引的损失计算
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# Adam优化器
optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)

def train_epoch(model, train_iter, optimizer):
    model.train()  # 设置模型为训练模式
    losses = 0  # 初始化损失值
    for idx, (src, tgt) in enumerate(train_iter):
        src = src.to(device)
        tgt = tgt.to(device)

        tgt_input = tgt[:-1, :]  # 目标输入序列（去掉最后一个元素）

        # 创建掩码
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

        # 模型前向传播
        logits = model(src, tgt_input, src_mask, tgt_mask,
                       src_padding_mask, tgt_padding_mask, src_padding_mask)

        optimizer.zero_grad()  # 清除梯度

        tgt_out = tgt[1:, :]  # 目标输出序列（去掉第一个元素）
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))  # 计算损失
        loss.backward()  # 反向传播

        optimizer.step()  # 优化步骤
        losses += loss.item()  # 累加损失值
    return losses / len(train_iter)  # 返回平均损失

def evaluate(model, val_iter):
    model.eval()  # 设置模型为评估模式
    losses = 0  # 初始化损失值
    for idx, (src, tgt) in enumerate(val_iter):
        src = src.to(device)
        tgt = tgt.to(device)

        tgt_input = tgt[:-1, :]  # 目标输入序列（去掉最后一个元素）

        # 创建掩码
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

        # 模型前向传播
        logits = model(src, tgt_input, src_mask, tgt_mask,
                       src_padding_mask, tgt_padding_mask, src_padding_mask)
        tgt_out = tgt[1:, :]  # 目标输出序列（去掉第一个元素）
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))  # 计算损失
        losses += loss.item()  # 累加损失值
    return losses / len(val_iter)  # 返回平均损失

2.6 开始训练

  最后，在准备好必要的类和函数之后，我们准备开始训练我们的模型。虽然不言而喻，但完成训练所需的时间可能会因计算能力、参数和数据集大小等许多因素而大不相同。

  当我使用JParaCrawl的完整句子列表训练模型时，每种语言大约有590万句子，使用单个NVIDIA GeForce RTX 3070 GPU，每个epoch大约需要5小时。

以下是代码： 
1
2
3
4
5
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS + 1), desc="Epochs"):
    start_time = time.time()  # 记录开始时间
    train_loss = train_epoch(transformer, train_iter, optimizer)  # 一个训练轮次
    end_time = time.time()  # 记录结束时间
    print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Epoch time = {(end_time - start_time):.3f}s")  # 打印损失和耗时

训练过程-by GPU

训练过程-by GPU

2.7 尝试使用训练好的模型翻译日语句子

  首先，我们创建一些函数来翻译新的句子，包括以下步骤：获取日语句子、分词、转换为张量、进行推断，然后将结果解码回英语句子。 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(device)  # 将源序列移动到设备上
    src_mask = src_mask.to(device)  # 将源序列掩码移动到设备上
    memory = model.encode(src, src_mask)  # 编码源序列，得到记忆向量
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device)  # 初始化目标序列，填充起始符号

    for i in range(max_len - 1):
        memory = memory.to(device)
        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(device).type(torch.bool)  # 生成记忆向量的掩码
        tgt_mask = (generate_square_subsequent_mask(ys.size(0)).type(torch.bool)).to(device)  # 生成目标序列的掩码
        out = model.decode(ys, memory, tgt_mask)  # 解码记忆向量，生成输出
        out = out.transpose(0, 1)
        prob = model.generator(out[:, -1])  # 通过生成器获取输出概率
        _, next_word = torch.max(prob, dim=1)  # 选取概率最大的词
        next_word = next_word.item()
        ys = torch.cat([ys, torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)  # 将选取的词添加到目标序列中
        if next_word == EOS_IDX:  # 如果选取的词是结束符号，停止解码
            break
    return ys

def translate(model, src, src_vocab, tgt_vocab, src_tokenizer):
    model.eval()  # 设置模型为评估模式
    # 将源序列分词，并添加起始符号和结束符号
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)] + [EOS_IDX]
    num_tokens = len(tokens)
    src = torch.LongTensor(tokens).reshape(num_tokens, 1)  # 将分词后的源序列转换为张量
    src_mask = torch.zeros(num_tokens, num_tokens).type(torch.bool)  # 生成源序列的掩码
    tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()  # 贪心解码生成目标序列
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")  # 返回翻译后的目标序列
  然后，我们只需调用translate函数并传递所需的参数即可 
1
2
3
4
5
6
7
8
# 翻译给定的日文句子
translation = translate(
    transformer,  # 使用的模型
    "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)",  # 要翻译的日文句子
    ja_vocab,  # 日文词汇表
    en_vocab,  # 英文词汇表
    ja_tokenizer  # 日文分词器
)

翻译日语

翻译日语

  原句使用DeepL的翻译结果为“HS 编码 8515 焊接、钎焊或熔接设备（电气设备（包括电加热燃气设备）”

1
trainen.pop(5)

结果为：'Chinese HS Code Harmonized Code System < HS编码 8515 : 电气(包括电热气体)、激光、其他光、光子束、超声波、电子束、磁脉冲或等离子弧焊接机器及装置,不论是否 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...'

1
trainja.pop(5)

结果为：'Japanese HS Code Harmonized Code System < HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)、レーザーその他の光子ビーム式、超音波式、電子ビーム式、 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...'

2.8 保存词汇对象和训练好的模型

  在训练完成后，我们将首先使用Pickle保存词汇对象（en_vocab和ja_vocab）。 

1
2
3
4
5
6
7
8
9
10
11
12
13
import pickle

# 存储英文词汇表数据
file = open('en_vocab.pkl', 'wb')
# 将英文词汇表数据序列化并存储到文件中
pickle.dump(en_vocab, file)
file.close()

# 存储日文词汇表数据
file = open('ja_vocab.pkl', 'wb')
# 将日文词汇表数据序列化并存储到文件中
pickle.dump(ja_vocab, file)
file.close()
  最后，我们还可以使用PyTorch的保存和加载功能保存模型以供日后使用。通常，有两种保存模型的方法，具体取决于我们以后想要如何使用它们。第一种方法是仅用于推断，我们可以稍后加载模型并用它来进行日语到英语的翻译。

1
2
# 保存模型的状态字典，以便进行推理
torch.save(transformer.state_dict(), 'inference_model')

  第二种方法也是用于推断，但适用于我们稍后想加载模型并继续训练的情况。 

1
2
3
4
5
6
7
# 保存模型和检查点，以便以后恢复训练
torch.save({
    'epoch': NUM_EPOCHS,  # 当前训练轮数
    'model_state_dict': transformer.state_dict(),  # 模型的状态字典
    'optimizer_state_dict': optimizer.state_dict(),  # 优化器的状态字典
    'loss': train_loss,  # 当前训练损失值
}, 'model_checkpoint.tar')  # 路径和名称

保存模型后的目录

保存模型后的目录
基于Transformer实现机器翻译
http://example.com/2024/06/22/基于Transformer实现机器翻译（日译中）/
作者
冷酷包
发布于
2024年6月22日
许可协议