基于 Transformer 和 PyTorch 的日中机器翻译模型

目录

一、获取语料库

1.1 导入需要的包

1.2 获取平行语料库

二、数据准备与加载

2.1 准备分词器

2.2 构建 TorchText Vocab 对象

2.3 创建 DataLoader 对象

三、Seq2Seq 翻译

四、模型训练

五、模型使用

参考文献

---

        在自然语言处理领域，机器翻译是一项具有挑战性的任务，它要求模型不仅要理解源语言文本的含义，还要能够准确地将其转换成目标语言。近年来，Transformer模型凭借其在处理序列数据方面的优势，已经成为机器翻译领域的主流架构之一。本实验旨在探索基于Transformer和PyTorch框架的日中机器翻译模型，通过实现Seq2SeqTransformer模型，我们能够深入理解Transformer的核心原理及其在机器翻译中的应用。

一、获取语料库

1.1 导入需要的包

        首先，确保系统中已经安装了以下包。

import math                          # 导入math库，用于访问数学函数和常数
import torchtext                     # 导入torchtext，用于文本处理
import torch                         # 导入PyTorch库，用于深度学习
import torch.nn as nn                # 从PyTorch导入nn模块，包含构建神经网络所需的类
from torch import Tensor             # 从PyTorch导入Tensor类，用于表示多维数组
from torch.nn.utils.rnn import pad_sequence # 导入pad_sequence函数，用于填充序列
from torch.utils.data import DataLoader # 导入DataLoader类，用于加载数据集
from collections import Counter      # 从collections模块导入Counter类，用于计数
from torchtext.vocab import Vocab    # 从torchtext导入Vocab类，用于构建词汇表
from torch.nn import TransformerEncoder, TransformerDecoder, \
    TransformerEncoderLayer, TransformerDecoderLayer # 从torch.nn导入Transformer模型的相关层
import io                            # 导入io库，用于处理输入输出流
import time                          # 导入time库，用于时间相关操作
import pandas as pd                  # 导入pandas库，用于数据分析和操作DataFrame对象
import numpy as np                   # 导入numpy库，用于数值计算
import pickle                        # 导入pickle库，用于序列化和反序列化Python对象
import tqdm                          # 导入tqdm库，用于显示进度条
import sentencepiece as spm          # 导入sentencepiece库，用于文本的分词处理
torch.manual_seed(0)                 # 设置PyTorch的随机种子为0，以确保结果的可重复性
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 设置设备为GPU（如果可用）或CPU

1.2 获取平行语料库

        这个实验中，我们将使用从JParaCrawl下载的日英平行语料库。JParaCrawl的数据是通过自动化的方式从互联网上抓取的，它代表了真实世界的语言使用情况。这些数据可以用来训练机器翻译模型，以提高从一种语言到另一种语言的翻译质量。

        导入了所有日语及其对应的英语文本之后，我删除了数据集中的最后一条数据，因为它存在一个缺失值。总的来说，在trainen和trainja中，句子的数量都是5,973,071条。然而，为了学习目的，通常建议在一次性使用所有数据之前，先对数据进行抽样并验证一切是否正常，这样可以节省时间。

# 使用pandas的read_csv函数读取文件，使用制表符'\t'作为分隔符，指定使用Python引擎来解析，且没有表头
df = pd.read_csv('./zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 从DataFrame的第三列（索引从0开始）提取数据，并将其转换为列表
trainen = df[2].values.tolist()#[:10000]

# 从DataFrame的第四列提取数据，并将其转换为列表
trainja = df[3].values.tolist()#[:10000]

# 查看索引为500的元素，以确保数据的正常
print(trainen[500])
print(trainja[500])

输出语料库中索引为500的元素

二、数据准备与加载

2.1 准备分词器

        在处理JParaCrawl语料库时，我们采用了其提供的分词器。这个分词器基于SentencePiece算法，能够为日语和英语文本生成合适的分词单元。SentencePiece是一个开源的、基于BPE（Byte Pair Encoding）的分词库，它能够处理多种语言，并且可以自动学习最优的分词方式。通过这种方式，我们能够更精确地对日英双语文本进行分词处理，为后面的机器翻译做准备。

# 创建一个英文分词器实例，使用预训练的SentencePiece模型
# 'spm.en.nopretok.model' 是模型文件的路径，这个模型用于英文分词
# 'nopretok' 表示模型不使用预标记的分词，即模型将直接从原始文本中生成标记
en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')

# 创建一个日文分词器实例，使用预训练的SentencePiece模型
# 'spm.ja.nopretok.model' 是模型文件的路径，这个模型用于日文分词
# 与英文分词器类似，这里也不使用预标记的分词
ja_tokenizer = spm.SentencePieceProcessor(model_file='spm.ja.nopretok.model')

# 使用英文分词器对给定的英文句子进行编码
# en_tokenizer.encode() 方法将句子转换为一系列分词后的标记
en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type='str')

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

         在分词器加载完成后，使用代码测试的结果：

2.2 构建 TorchText Vocab 对象

        定义 build_vocab 的函数，它的作用是根据一组句子和分词器来构建一个词汇表。这个函数首先使用 Counter 类来计数句子中的不同标记，Counter 是一个字典子类，用于计数可哈希对象。函数遍历传入的 sentences 列表，对每一句使用 tokenizer 进行编码，将句子转换成标记序列，并将这些标记更新到 counter 中。

  tokenizer.encode 方法将句子转换为一系列标记，并且 out_type=str 参数指定输出类型为字符串。然后，使用 counter.update 方法将这些标记添加到计数器中，从而统计每个标记出现的次数。

        接下来，函数使用 counter 来构建词汇表。这里使用了 Vocab 构造函数，它接受一个计数器对象，并可以添加一些特殊标记，如未知标记 <unk>、填充标记 <pad>、开始标记 <bos> 和结束标记 <eos>。这些特殊标记在自然语言处理中有着重要的作用，例如处理未知词汇、填充序列以统一长度、标记句子的开始和结束。

        最后，函数返回构建好的词汇表。在代码的后半部分，分别使用 build_vocab 函数和日语、英语的分词器 ja_tokenizer、en_tokenizer 来构建日语和英语的词汇表，分别存储在 ja_vocab 和 en_vocab 变量中。这样，每个语言的文本在进行进一步的自然语言处理任务之前，都会有一个准备好的词汇表，这有助于标准化和优化模型的训练过程。

# 定义一个函数 build_vocab，用于根据一组句子和分词器构建词汇表
def build_vocab(sentences, tokenizer):
    # 使用 Counter 对象来计数句子中的不同标记
    counter = Counter()
    
    # 遍历所有句子
    for sentence in sentences:
        # 使用分词器对句子进行编码，并以字符串形式返回标记序列
        # counter.update() 方法将更新计数器，包括句子中的所有标记
        counter.update(tokenizer.encode(sentence, out_type=str))
    
    # 使用计数器构建一个词汇表
    # Vocab 构造函数接受计数器，并可以指定特殊标记
    # 特殊标记通常包括未知标记 '<unk>', 填充标记 '<pad>', 开始标记 '<bos>' 和结束标记 '<eos>'
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])

# 使用 build_vocab 函数和 ja_tokenizer 构建日语词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)

# 使用 build_vocab 函数和 en_tokenizer 构建英语词汇表
en_vocab = build_vocab(trainen, en_tokenizer)

        定义 data_process 的函数，其主要作用是处理原始的日语和英语句子数据，将它们转换成适合机器学习模型输入的格式。函数首先初始化一个空列表 data，用于存储处理后的数据。

        接下来，函数使用 zip 函数将日语列表 ja 和英语列表 en 配对组合，并遍历这些配对。对于每一对原始日语句子 raw_ja 和原始英语句子 raw_en，函数执行以下步骤：

1. 使用 ja_tokenizer 对日语句子进行分词，并将分词结果转换为字符串列表，同时去除字符串末尾的换行符。
2. 将分词后的字符串列表 ja_tokens 中的每个token转换为对应 ja_vocab 词汇表中的索引，并创建一个 PyTorch 张量 ja_tensor_，其数据类型为 torch.long，这表示张量中的元素是长整型。
3. 以相同的方式处理英语句子，使用 en_tokenizer 分词，并将token转换为 en_vocab 词汇表中的索引，创建对应的 PyTorch 张量 en_tensor_。

        处理后的日语和英语张量被组成一个元组 (ja_tensor_, en_tensor_)，然后这个元组被添加到 data 列表中。这样，每对原始句子都被转换成了它们的数值表示形式，并且存储在列表中。

        最后，函数返回这个包含处理后数据的列表 data。通过调用 data_process 函数并传入训练数据 trainja 和 trainen，得到了处理后的训练数据集 train_data。这个数据集现在可以被用来创建 DataLoader 对象，进而在模型训练过程中进行批量加载和处理。

def data_process(ja, en):
    # 初始化一个空列表，用于存储处理后的数据
    data = []

    # 使用zip函数将ja和en两个列表组合起来，然后遍历这个组合
    for (raw_ja, raw_en) in zip(ja, en):
        # 使用ja_tokenizer对raw_ja进行编码，并将编码后的字符串列表中的每个token转换为ja_vocab中的索引
        # rstrip("\n")用于去除字符串末尾的换行符
        # out_type=str指定输出类型为字符串
        ja_tokens = ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)
        # 将ja_tokens中的每个token转换为对应的索引，并创建一个PyTorch张量（tensor），数据类型为long
        ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokens], dtype=torch.long)

        # 同理，对raw_en进行相同的处理
        en_tokens = en_tokenizer.encode(raw_en.rstrip("\n"), out_type=str)
        en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokens], dtype=torch.long)

        # 将处理后的ja_tensor_和en_tensor_组成一个元组，并添加到data列表中
        data.append((ja_tensor_, en_tensor_))

    # 返回处理后的数据列表
    return data

# 调用data_process函数，传入训练数据trainja和trainen，得到处理后的训练数据集train_data
train_data = data_process(trainja, trainen)

2.3 创建 DataLoader 对象

        下面，我们定义如何生成训练批次数据，并使用这些数据创建了一个迭代器，以便在模型训练过程中使用。

        首先，定义了三个常量：BATCH_SIZE 设置为 16，表示每个批次包含 16 个样本；PAD_IDX 是日语词汇表中 <pad> 标记的索引，用于填充序列以确保它们具有相同的长度；BOS_IDX 和 EOS_IDX 分别是 <bos> 和 <eos> 标记的索引，分别用于标记句子的开始和结束。

        接着，定义了一个名为 generate_batch 的函数，它接收一个批次的数据 data_batch 作为输入。这个函数的作用是将批次中的每个日语和英语句子处理成统一格式的张量。对于每对 (ja_item, en_item)：

• 使用 torch.cat 函数将 <bos> 标记的索引添加到句子的开始，将 <eos> 标记的索引添加到句子的结束，并将这些索引转换为张量。
• 将处理后的日语和英语句子分别添加到 ja_batch 和 en_batch 列表中。

        然后，使用 pad_sequence 函数对 ja_batch 和 en_batch 列表中的张量进行填充，确保所有序列具有相同的长度。padding_value=PAD_IDX 参数指定使用 <pad> 标记的索引作为填充值。

        最后，使用 DataLoader 类创建了一个名为 train_iter 的迭代器。这个迭代器接收处理后的训练数据集 train_data，设置批次大小为 BATCH_SIZE，并且设置 shuffle=True 以在每个epoch开始时打乱数据。collate_fn=generate_batch 参数指定了如何将多个数据样本合并为一个批次，即使用 generate_batch 函数。

BATCH_SIZE = 16
PAD_IDX = ja_vocab['<pad>']
BOS_IDX = ja_vocab['<bos>']
EOS_IDX = ja_vocab['<eos>']
def generate_batch(data_batch):
    ja_batch, en_batch = [], []
    for (ja_item, en_item) in data_batch:
        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
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

        通过这种方式，train_iter 迭代器能够在训练过程中按批次提供处理好的数据，使得模型可以逐批次地学习，从而提高训练效率和稳定性。 

三、Seq2Seq 翻译

        下面，我们定义一个基于Transformer的序列到序列（Seq2Seq）模型，用于处理序列数据的翻译或转换任务。模型由编码器（encoder）和解码器（decoder）组成，它们共享相同的维度和一些层的参数。

        首先，代码从PyTorch的torch.nn模块导入了必要的Transformer层和组件，包括TransformerEncoder、TransformerDecoder、TransformerEncoderLayer和TransformerDecoderLayer。

  Seq2SeqTransformer类继承自nn.Module，是一个神经网络模块。在初始化方法__init__中，定义了模型的主要组成部分：

• encoder_layer和decoder_layer：分别创建Transformer编码器和解码器层。
• transformer_encoder和transformer_decoder：构建编码器和解码器，它们分别由多个编码器层和解码器层堆叠而成。
• generator：一个线性层，用于将解码器的输出转换为目标词汇表大小的概率分布。
• src_tok_emb和tgt_tok_emb：分别为源语言和目标语言的词嵌入层。
• positional_encoding：位置编码，为模型提供序列中单词的位置信息。

  forward方法定义了模型的前向传播过程：

• 使用词嵌入层和位置编码处理源语言和目标语言的输入序列。
• 使用编码器处理源语言的嵌入序列。
• 使用解码器处理目标语言的嵌入序列和编码器的输出。
• 最后，使用generator层生成最终的输出。

  encode方法用于执行源语言序列的编码过程，返回编码器的输出。 decode方法用于执行目标语言序列的解码过程，给定目标语言的序列、编码器的输出和解码器的掩码。

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__()
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        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))
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        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)

        整个模型的设计利用了Transformer架构的优势，包括自注意力机制和位置编码，以实现高效的序列转换任务。

        下面将定义两个类，它们都是PyTorch的nn.Module的子类，用于为Transformer模型提供位置编码和词嵌入。

1. PositionalEncoding类：

• 构造函数接收三个参数：emb_size表示嵌入的维度，dropout表示dropout比率，maxlen表示最大序列长度。
• 通过计算一个分母den，用于对不同维度的正弦和余弦函数的频率进行缩放。
• 使用torch.arange生成位置向量pos，然后根据位置和den计算正弦和余弦值，创建位置编码矩阵pos_embedding。
• 通过unsqueeze操作，为位置编码添加一个维度，使其可以与词嵌入相加。
• 定义了一个nn.Dropout层，用于在训练时随机丢弃一些位置编码，以提高模型的泛化能力。
• 使用register_buffer注册位置编码为模型的缓冲区，这意味着它不会被优化器更新。

forward方法：

• 接收词嵌入token_embedding作为输入。
• 将dropout应用于词嵌入。
• 将位置编码添加到词嵌入中，确保位置编码的长度与词嵌入的序列长度一致。

2. TokenEmbedding类：

• 构造函数接收两个参数：vocab_size表示词汇表的大小，emb_size表示嵌入的维度。
• 定义了一个nn.Embedding层，用于将词汇表中的每个词映射到一个嵌入向量。
• 将嵌入向量乘以emb_size的平方根，这有助于稳定训练过程。

forward方法：

• 接收一个包含词索引的张量tokens作为输入。
• 使用self.embedding获取词嵌入，并将输入转换为长整型，因为nn.Embedding需要长整型索引。
• 返回的词嵌入将被传递给位置编码模块。

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)
        pos_embedding = pos_embedding.unsqueeze(-2)

        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        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):
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

        下面定义了两个函数，它们用于在Transformer模型中生成不同类型的掩码。

generate_square_subsequent_mask函数：

• 接收一个参数sz，表示矩阵的大小。
• 使用torch.triu创建一个上三角矩阵，其中的对角线和下三角部分为1，其余部分为0。
• 将结果转换为浮点数，并使用masked_fill将值为0的位置替换为负无穷大（float('-inf')），将值为1的位置替换为0.0。这样做是为了在后续的计算中，通过softmax函数将上三角部分的值压缩到接近0，而下三角部分的值保持不变。

create_mask函数：

• 接收两个参数src和tgt，分别代表源序列和目标序列。
• 首先计算源序列和目标序列的长度。
• 调用generate_square_subsequent_mask函数生成目标序列的后续掩码tgt_mask，这个掩码用于防止解码器在生成序列时看到未来的信息（mask future tokens）。
• 创建源序列的掩码src_mask，初始化为全0的布尔矩阵，表示不进行任何掩码操作。
• 生成源序列的填充掩码src_padding_mask和目标序列的填充掩码tgt_padding_mask，通过检查序列中的填充索引（PAD_IDX），将填充位置标记为True，其他位置标记为False。这些掩码用于在计算损失时忽略填充位置。

def generate_square_subsequent_mask(sz):
    # 创建一个上三角矩阵，并转换为浮点数
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将上三角部分的值替换为负无穷大，下三角部分的值替换为0
    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_mask = generate_square_subsequent_mask(tgt_seq_len)
    # 初始化源序列掩码为全0的布尔矩阵
    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

        这些掩码在Transformer模型中至关重要，它们确保解码器在生成每个输出时刻的词时，只能看到当前和之前的词，同时在计算损失时忽略填充的部分。 

        接下来，我们将配置和训练一个基于Transformer的序列到序列模型，包括模型初始化、参数初始化、设备分配、损失函数和优化器的定义，以及训练和评估过程的实现。

        在构建Seq2SeqTransformer模型的过程中，首先定义了一系列关键的超参数，包括源语言和目标语言的词汇表大小SRC_VOCAB_SIZE和TGT_VOCAB_SIZE，嵌入层的维度EMB_SIZE，多头注意力机制的头数NHEAD，前馈网络隐藏层的维度FFN_HID_DIM，批次大小BATCH_SIZE，以及编码器和解码器的层数NUM_ENCODER_LAYERS和NUM_DECODER_LAYERS，还设置了训练轮数NUM_EPOCHS。接着，基于这些超参数配置，创建了transformer模型实例。为了促进训练过程中的梯度流动，模型的所有参数采用了Xavier均匀初始化方法。随后，模型被移动到指定的设备上，可能是CPU或GPU，以便进行高效计算。

        在训练准备阶段，定义了一个交叉熵损失函数loss_fn，特别地，它能够忽略填充索引PAD_IDX，以避免在损失计算时对填充的部分进行惩罚。此外，创建了一个Adam优化器optimizer，它将用于后续模型训练过程中的参数更新。

        训练过程通过train_epoch函数实现，该函数将模型置于训练模式，并在每个epoch中循环遍历训练迭代器train_iter。在每个批次的处理中，源语言和目标语言的数据被送至设备，目标序列生成相应的输入（排除序列的最后一个时间步），同时创建了模型所需的掩码。模型据此产生预测输出logits，随后进行梯度清零、损失计算、反向传播以及参数更新，最终返回该epoch的平均损失。

        模型评估则通过evaluate函数完成，在评估模式下，模型对验证迭代器val_iter上的数据进行处理，计算损失，但不执行反向传播和参数更新，最终得到验证集的平均损失，从而评估模型的性能。

# 设置源语言和目标语言词汇表的大小
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)

# 将模型移动到指定的设备（CPU或GPU）
transformer = transformer.to(device)

# 定义损失函数，使用交叉熵损失，并忽略填充索引PAD_IDX
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
)

# 定义训练函数，用于训练模型一个epoch
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()

    # 返回该epoch的平均损失
    return losses / len(train_iter)

# 定义评估函数，用于评估模型在验证集上的性能
def evaluate(model, val_iter):
    model.eval()  # 设置模型为评估模式
    losses = 0  # 初始化损失累计变量
    # 遍历验证迭代器
    for idx, (src, tgt) in enumerate(valid_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)

四、模型训练

        迭代训练模型指定的次数，即NUM_EPOCHS轮。在每次训练（即每个epoch）中，执行以下操作：

1. 使用tqdm.tqdm包装range函数生成的序列，从1到NUM_EPOCHS+1，这样可以在控制台显示一个进度条，表示训练进程和每个epoch的完成状态。

2. 记录当前时间start_time，这是为了计算训练一个epoch所需的时间。

3. 调用train_epoch函数并传入模型transformer、训练迭代器train_iter和优化器optimizer，以此执行一个epoch的训练，并获取该epoch的训练损失train_loss。

4. 记录训练完成后的时间end_time。

5. 打印出当前epoch的编号、训练损失和该epoch的训练时间。

import time
from tqdm import tqdm

# 训练模型的循环
for epoch in tqdm(range(1, NUM_EPOCHS+1)):  # 使用tqdm显示进度条
    start_time = time.time()  # 记录epoch开始时间
    
    # 训练一个epoch并获取损失
    train_loss = train_epoch(transformer, train_iter, optimizer)
    
    end_time = time.time()  # 记录epoch结束时间
    # 打印epoch编号、训练损失和训练时间
    print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
           f"Epoch time = {(end_time - start_time):.3f}s"))

五、模型使用

        尝试使用上面训练好的模型来实现日语句子的翻译。

        定义下面两个函数用于使用训练好的Seq2SeqTransformer模型进行解码和翻译。

  greedy_decode函数实现了一种简单的解码策略，它接收一个预训练的模型、源语言文本的表示、相应的掩码、允许的最大解码长度以及解码开始的起始符号。函数首先将源语言数据准备到合适的设备上，并利用模型的编码器部分来获取编码后的记忆表示。接着，初始化目标序列，仅包含起始符号，并在设备上进行处理。在每次迭代中，函数生成目标掩码以防止解码时信息泄露，然后利用解码器和当前的目标序列生成下一个词的概率分布。之后，选择概率最高的词作为序列的下一个元素，直到生成结束符号或达到最大长度。最后，函数返回解码得到的完整目标序列。

  translate函数则是将greedy_decode函数应用于实际翻译任务中的一个包装器。它首先将模型设置为评估模式，使用分词器将源语言文本转换为词索引，并添加必要的开始和结束符号。然后，构建源语言数据的张量和掩码，调用greedy_decode函数进行解码。解码得到的词索引序列经过转换和格式化，去除特定的标记，并连接成人类可读的目标语言文本。这样，translate函数提供了一个从源语言文本到目标语言文本的直接翻译接口。

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    # 将源语言数据和掩码移动到指定的设备（CPU或GPU）
    src = src.to(device)
    src_mask = src_mask.to(device)
    # 使用模型的编码器对源数据进行编码，得到源记忆memory
    memory = model.encode(src, src_mask)
    
    # 初始化目标序列ys，填充为起始符号，并移动到设备上
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device)
    
    # 进行循环直到达到最大长度或生成结束符号
    for i in range(max_len-1):
        # 确保memory在正确的设备上
        memory = memory.to(device)
        # 创建memory_mask以屏蔽源记忆与目标序列之间的交互
        memory_mask = torch.zeros(ys.shape[0], memory.shape[0]).to(device).type(torch.bool)
        # 生成tgt_mask以防止解码器在未来步骤中看到未来的信息
        tgt_mask = (generate_square_subsequent_mask(ys.size(0)).type(torch.bool)).to(device)
        
        # 使用模型的解码器和目标掩码对目标序列进行解码，得到输出out
        out = model.decode(ys, memory, tgt_mask)
        # 转置输出out以匹配生成器层的输入要求
        out = out.transpose(0, 1)
        # 通过模型的生成器层获取下一个词的概率分布
        prob = model.generator(out[:, -1])
        # 从概率分布中选择概率最高的词作为下一个词
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.item()  # 转换为Python标量
        # 将新词添加到目标序列ys中
        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]
    
    # 将词索引转换为张量src，并创建相应的源掩码src_mask
    num_tokens = len(tokens)
    src = torch.LongTensor(tokens).reshape(num_tokens, 1)
    src_mask = torch.zeros(num_tokens, num_tokens).type(torch.bool)
    
    # 调用greedy_decode函数进行解码
    # 设置最大长度为源序列长度加5，起始符号为BOS_IDX
    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>", "")

        这两个函数共同实现了从源语言到目标语言的翻译过程，首先对源文本进行编码，然后逐词解码并生成目标文本。greedy_decode函数采用了贪婪策略，每次都选择概率最高的词，直到生成结束符号或达到最大长度。translate函数则负责准备数据、调用解码函数和后处理，最终输出翻译结果。

        最后，调用translate函数并传入所需的参数，将模型应用在一个日语句子的翻译上。

translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)

参考文献

https://proceedings.neurips.cc/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf