基于pytorch的Transformer日译中模型

Transformer 是一种基于注意力机制的深度学习模型，最初由Vaswani等人在2017年提出，主要用于处理序列到序列（Sequence-to-Sequence，Seq2Seq）的任务，如机器翻译和文本生成。相比于传统的循环神经网络（RNN）和长短时记忆网络（LSTM），Transformer 引入了自注意力机制（Self-Attention）来捕捉输入序列中各个位置之间的依赖关系，同时避免了传统序列模型中存在的顺序依赖问题。下面是transformer整体结构:

主要组成部分：位置编码（Positional Encoding）、多头注意力层（Multi-Head Attention）、带掩码的多头注意力层（ Masked Multi-Head Attention）、前馈神经网络层（Feed Forward）、残差连接和层归一化层（Add & Norm）

接下来将逐一介绍各个主要组成部分的构成和实现。

1.2位置编码（Positional Encoding)

Transformer 使用位置编码来为序列中每个位置添加信息，以表示它们在序列中的相对或绝对位置。

实际上，在传统的RNN中的序列是携带着位置信息的，因为RNN是顺序执行的。因此，对传统的RNN而言，没有必要额外添加位置编码层。

transformer中的自注意力机制在计算的过程中进行完全并行化，这一方面大大提高了训练和推理的速度，但带来的问题就是不能携带原先序列中的位置信息。为此设计出了位置编码层，为序列添加上位置信息，表示序列中的相对或绝对位置。

位置编码通常是通过正弦和余弦函数的组合来实现的，这样可以使得模型能够理解序列中不同位置的顺序信息。具体的实现算法和过程这里不再展开，详情请参考【Transformer系列】深入浅出理解Positional Encoding位置编码-CSDN博客

1.3多头注意力（Multi-Head Attention）

在介绍多头注意力之前，让我们先来看看其基础——自注意力机制（Self-Attention Mechanism）。

自注意力机制是Transformer模型的核心组件之一，它能够捕捉输入序列中各个位置之间的依赖关系，从而在处理序列数据时具有更高的并行度和灵活性。以下是自注意力机制的关键概念和计算步骤：

关键概念：

查询（Query）：用于查找相关信息的向量。
键（Key）：表示输入序列中每个位置的信息。
值（Value）：包含与每个键对应的实际信息。

计算步骤：

1.线性变换：

将输入序列中的每个元素通过三组权重矩阵分别映射为查询（Q）、键（K）和值（V）向量：

$Q=XW_{Q},K=XW_{K},V=XW_{V}$

其中， $X$ 是输入序列， $W_{Q},W_{K},W_{V}$ 是可学习的权重矩阵。

2.计算注意力得分：

通过点积计算每个查询与所有键的相似度，然后除以一个缩放因子 $\sqrt{d_{k}}$ 以避免点积结果过大：

$Attention Score(Q,K)=\frac{QK^{T}}{\sqrt{d_k}}$

其中， $\sqrt{d_{k}}$ 是键向量的维度。

3.应用Softmax：

对注意力得分进行Softmax操作，将其转换为权重，表示每个位置对其他位置的关注程度：

$Attention Weights=softmax(\frac{QK^{T}}{\sqrt{d_{k}}})$

4.加权求和：

将注意力权重与值向量相乘，并对结果求和，得到每个位置的最终表示：

$Output=Attention Weights\cdot V$

如果我们选择用矩阵处理，可以直接将之前的第二至四步压缩到一个公式中一步到位获得最终的注意力结果，如下图所示：

自注意力机制通常会扩展为多头注意力机制，即将上述计算步骤在多个注意力头上独立执行。每个注意力头使用不同的权重矩阵，并生成不同的查询、键和值向量。最后，将所有注意力头的输出拼接起来，再通过一个线性变换得到最终的输出。

$Multi-Head Output=Concat(head_{1},head_{2},...,head_{H})W_{O}$

其中， $H$ 是注意力头的数量， $W_{O}$ 是输出的权重矩阵。

多头注意力机制（Multi-Head Attention）是Transformer模型的重要组成部分，相比单头注意力机制，它具有以下几个显著优势：

1. 捕捉多种特征和关系

每个注意力头独立地学习和捕捉不同的特征和关系。通过在多个头上进行注意力计算，模型可以在不同的子空间中并行地捕捉到输入序列中不同位置之间的多种依赖关系。这使得模型能够更好地理解复杂的上下文信息。

2. 增强表示能力

由于多头注意力机制能够从不同的角度处理和组合信息，因此它增强了模型的表示能力。每个头的关注点不同，它们在一起能够提供更加丰富和多样的表示，从而提高模型的整体表现。

3. 提高稳定性和鲁棒性

多头注意力机制通过将输入信息分散到多个头上，降低了模型对单一头部学习效果的依赖。这种方式增加了模型的稳定性和鲁棒性，即使某个头的注意力权重不理想，其他头仍然可以提供有价值的信息。

4. 更细粒度的注意力分布

单头注意力可能会过于集中在某些特定位置，而忽略其他有用的信息。多头注意力机制允许模型在不同头上关注不同的位置，从而提供更细粒度的注意力分布。这有助于更全面地理解输入序列的整体结构和语义。

5. 改善梯度流动

多头注意力机制通过并行计算多个注意力头，增加了梯度流动的路径数。这使得梯度更容易传播，从而改善了模型的训练效果，尤其在处理深层网络时表现更为显著。

6. 有效利用隐藏维度

多头注意力机制将隐藏维度划分为多个子空间，每个头在较小的子空间中进行注意力计算。这种划分可以更有效地利用隐藏维度，提升模型的计算效率和性能。

1.4掩码注意力（Mask Self_Attention)

掩码注意力机制在Transformer中起到了关键作用，特别是在序列到序列的生成任务中，如机器翻译。在解码器部分，未来信息的掩码（Future Masking）确保在预测当前词时模型不能看到之后的词，从而维持自回归属性。这种设计保证了模型的训练和推理过程符合语言的自然生成顺序。

mask操作的实现：在注意力得分的基础上加上上三角矩阵Mask，其主对角线以上所有元素均为负无穷。即：

$Z=softmax(\frac{QK^{T}}{\sqrt{d_{k}}}+Mask)\cdot V$

其中， $Mask_{i,j}=\left\{\begin{matrix} 0,i\geqslant j\\ -\infty ,i<j \end{matrix}\right.$

Mask矩阵中负无穷的位置即代表未来信息的位置。对于未来信息的位置上，注意力得分加上负无穷，经过softmax后必定得到的是0。这样一来，未来信息成功被隐藏。

1.5前馈神经网络层（Feed Forward）

编码器和解码器中都包含前馈神经网络层，这些层在每个位置独立地进行计算，并通过激活函数（通常是ReLU）增强模型的表示能力。它位于多头自注意力机制之后，并且在每个输入位置上独立地应用。前馈神经网络帮助对注意力机制输出的信息进行进一步的处理和变换。

前馈神经网络通常包含两个线性变换（全连接层）和一个激活函数。具体来说，它的结构如下：

输入层：接收来自多头自注意力机制的输出。
隐藏层：应用第一个线性变换，将输入从维度d_model映射到更高维度d_ff，应用非线性激活函数（通常是ReLU）。
输出层：应用第二个线性变换，将隐藏层的输出映射回原始维度d_model。

假设输入为 𝑥，那么前馈神经网络的计算过程可以表示为：

$FFN(x)=max(0,xW_{1}+b_{1})W_{2}+b_{2}$

其中：

$W_{1}$ 和 $b_{1}$ 是第一层的权重和偏置。
$W_{2}$ 和 $b_{2}$ 是第二层的权重和偏置。
max⁡(0,⋅) 表示ReLU激活函数。

1.6残差连接和层归一化（Residual Connection & Layer Normalization）

在Transformer模型中，残差连接和层归一化是两个关键组件，它们在每个编码器和解码器层中使用，用于改善模型的训练稳定性和效果。

残差连接（Residual Connection）

残差连接通过在每个子层（子网络）的输入和输出之间添加跳跃连接来实现。这种连接方式有助于缓解梯度消失问题，并加速训练过程。残差连接的主要思想是通过直接将输入信息传递到输出，从而使得子层只需要学习输入与输出之间的残差。

公式表示为：

$output=input+sublayer(input)$

这种结构最早由He等人在ResNet（Residual Network）中提出，并且被广泛应用于各种神经网络结构中。

层归一化（Layer Normalization）

层归一化在每个子层后面使用，目的是标准化每个样本的特征，从而使得输入数据的分布更为稳定。这种标准化操作有助于加速模型的收敛，并改善模型的训练性能。

公式表示为：

$LayerNorm(x)=\frac{x-\mu }{\sigma }\cdot \gamma +\beta$

其中：

𝜇是输入的均值。
𝜎 是输入的标准差。
𝛾 和 𝛽是可学习的参数，用于恢复归一化前的缩放和平移。

残差连接和层归一化的组合

在Transformer中，残差连接和层归一化通常结合使用，以确保在每个子层操作后，网络的输入和输出保持稳定的分布。这种组合结构如下：

输入通过子层（如多头注意力或前馈神经网络）。
子层输出加上残差连接（即输入加上子层的输出）。
加和后的结果通过层归一化。

这种结构的公式表示为： $output=LayerNorm(input+sublayer(input))$

残差连接通过直接传递输入信息来缓解梯度消失问题，并加速训练过程。层归一化通过标准化每个样本的特征来确保输入数据的分布稳定，从而加速模型的收敛。这两者的结合使用，确保了Transformer在处理复杂任务时的稳定性和效率。

二、技术实现

2.1环境搭建

确保实验过程的可视化和代码模块化，本实验在Visual Studio Code中的jupyter notebook环境下进行。

导入后续代码中所需要的库：

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)

简单说一下我的版本：

torchtext==0.6.0

torch==2.3.1

torchvision==0.18.0

tqdm==4.65.0

sentencepiece==0.2.0

检查电脑配置，选择在GPU下运行代码，如果没有的话则为CPU：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(torch.cuda.get_device_name(0)) 
device

我是在GPU下运行的，上面代码的输出为：

NVIDIA GeForce RTX 3050 Laptop GPU

device(type='cuda')

2.2数据读取与预处理

首先，我们选取中日语料（from JParaCrawl）来进行本次实验。

# 使用pandas库读取一个制表符分隔的文件，并将其内容存储到一个DataFrame中
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
# 从DataFrame的第三列（即索引为2的列）提取数据，并转换为列表，存储到trainen变量中
trainen = df[2].values.tolist()#[:10000]
# 从DataFrame的第四列（即索引为3的列）提取数据，并转换为列表，存储到trainja变量中
trainja = df[3].values.tolist()#[:10000]
# trainen.pop(5972)
# trainja.pop(5972)

以下是数据集中包含的一个句子的示例：

# 打印元素
print(trainen[500])
print(trainja[500])

测试结果如下：

Chinese HS Code Harmonized Code System < HS编码 2905 无环醇及其卤化、磺化、硝化或亚硝化衍生物 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ... Japanese HS Code Harmonized Code System < HSコード 2905 非環式アルコール並びにそのハロゲン化誘導体、スルホン化誘導体、ニトロ化誘導体及びニトロソ化誘導体 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...

我们还可以使用不同的平行数据集来跟随本文，只需确保我们能够将数据处理成上面展示的两个字符串列表，其中包含日语和英语句子。

不像英语或其他字母语言，日语句子中没有空格来分隔单词。我们可以使用由JParaCrawl提供的分词器，它们分别用于日语和英语，这些分词器是使用SentencePiece创建的。

# 初始化英文分词器，使用 SentencePiece 模型 'spm.en.nopretok.model'
en_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')
# 初始化日文分词器，使用 SentencePiece 模型 'spm.ja.nopretok.model'
ja_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.ja.nopretok.model')

加载分词器后，可以通过执行以下代码来测试它们的功能，例如：

# 使用英文分词器对给定的英文句子进行编码，输出类型为字符串列表
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)

测试结果如下：

['▁All', '▁residents', '▁aged', '▁20', '▁to', '▁59', '▁years', '▁who', '▁live', '▁in', '▁Japan', '▁must', '▁enroll', '▁in', '▁public', '▁pension', '▁system', '.']

['▁', '年', '金', '▁日本', 'に住んでいる', '20', '歳', '~', '60', '歳の', '全ての', '人は', '、', '公的', '年', '金', '制度', 'に', '加入', 'しなければなりません', '。']

然后，我们用分词后的文本去创建输入输出的词汇表（词汇表不要忘记特殊字符的添加），利用词汇表完成整个数据集的映射处理，得到张量形式的训练数据。

# 定义一个构建词汇表的函数，接受句子列表和分词器作为参数
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)

# 定义一个数据处理函数，接受日文和英文列表作为输入
def data_process(ja, en):
  data = []
  # 使用zip函数同时遍历日文和英文列表
  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)

之后，为了便于模型训练，我们将数据加上开始结束索引，转换为批次的形式，并用Dataloader对象储存。

# 定义批量大小
BATCH_SIZE = 8
# 获取填充符号的索引
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:
    # 在日文句子的开头和结尾添加起始符号（<bos>）和终止符号（<eos>），并连接为一个张量
    ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
    # 在英文句子的开头和结尾添加起始符号（<bos>）和终止符号（<eos>），并连接为一个张量
    en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
  # 对日文批次进行填充，使其长度保持一致，使用PAD_IDX作为填充值
  ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
  # 对英文批次进行填充，使其长度保持一致，使用PAD_IDX作为填充值
  en_batch = pad_sequence(en_batch, padding_value=PAD_IDX)
  return ja_batch, en_batch
# 使用DataLoader加载训练数据，每次迭代返回一个批次数据
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

2.3 模型搭建

我们按照论文中所描述的Transformer架构来搭建模型，其中编码器和解码器由位置编码、词嵌入、多头注意力机制、残差连接、层归一化和前向传播等模块组成，区别在于解码器的自注意力层会采用mask对未来信息进行遮掩，而且解码器需要处理编码器的输出来结合上下文。此处采用nn中模块化的编码器层和解码器层。

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))
        # 通过编码器层处理源语言嵌入，生成编码器输出
        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)

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))
        # 对位置嵌入矩阵的偶数列应用sin函数
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        # 对位置嵌入矩阵的奇数列应用cos函数
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        # 在最后一维增加一个维度
        pos_embedding = pos_embedding.unsqueeze(-2)

        # 初始化Dropout层
        self.dropout = nn.Dropout(dropout)
        # 将位置嵌入矩阵注册为模型的buffer，使其在模型保存和加载时被保留
        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):
        # 返回词嵌入，并对其进行缩放
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

下面是mask的生成机制，解码器的mask是下三角的-inf矩阵（其中-inf表示该位置权重为0，达到遮掩信息的目的），用来遮掩未来信息，而编码器的mask是全零矩阵（0表示对该位置权重不造成影响），不需要对信息进行遮掩。

def generate_square_subsequent_mask(sz):
    # 生成上三角矩阵，大小为(sz, sz)，对角线及其上方为1，其余为0
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将mask转换为浮点数，并将0位置填充为负无穷大，将1位置填充为0.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

定义模型的各项参数，准备开始训练。（这里我还是头铁了，没修改参数，给电脑跑了六个多小时）

# 定义模型的各项参数
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)

# 定义损失函数，忽略填充标记的损失
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(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)  # 返回平均损失

2.4模型训练

在准备好必要的类和函数之后，我们准备训练我们的模型。但完成训练所需的时间可能会因许多因素而异，例如计算能力、参数和数据集的大小。以下为代码：

# 使用 tqdm 显示训练进度条
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time()  # 记录开始时间
  # 调用 train_epoch 函数进行训练，并返回训练损失
  train_loss = train_epoch(transformer, train_iter, optimizer)
  end_time = time.time()  # 记录结束时间
  # 打印当前轮次的训练损失和所用时间
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
          f"Epoch time = {(end_time - start_time):.3f}s"))

训练过程展示：

2.5翻译测试

首先，我们创建用于翻译新句子的函数，包括获取日语句子、进行标记化、转换为张量、推断，然后将结果解码回一个英语句子。在预测过程中，实验采用较为简单的贪婪解码方法，大致过程与训练是相似的，选取解码器输出中概率分布最高的索引。

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)
    # 初始化目标句子，起始符号为start_symbol
    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()  # 设置模型为评估模式
    # 将源句子编码为词ID，并添加起始符和结束符
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)]+ [EOS_IDX]
    num_tokens = len(tokens)
    # 将词ID转换为张量并调整形状
    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()
    # 将目标句子的词ID转换为词，并去掉起始符和结束符
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

然后，我们只需调用翻译函数并传入必要的参数即可。

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

翻译结果：

' ▁H S ▁ 代码 ▁85 15 ▁ 用于焊接应用的焊接设备 ( 包括电热加热 ) 。 '

2.6保存词汇表和模型

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

import pickle
# open a file, where you want to store the data
file = open('en_vocab.pkl', 'wb')
# dump information to that file
pickle.dump(en_vocab, file)
file.close()
file = open('ja_vocab.pkl', 'wb')
pickle.dump(ja_vocab, file)
file.close()

最后，我们还可以使用PyTorch的保存和加载函数来保存模型以便将来使用。通常，根据我们想要将模型用于以后的目的，有两种保存模型的方式。第一种是仅用于推断，我们可以稍后加载模型并使用它来从日语翻译到英语。

# 保存用于推断的模型
torch.save(transformer.state_dict(), 'inference_model')

第二种方法也是用于推断，但是它还用于当我们想要稍后加载模型并希望恢复训练时。

# 将模型和检查点保存以便稍后恢复训练
torch.save({
    'epoch': NUM_EPOCHS,                    # 保存当前训练的轮次数
    'model_state_dict': transformer.state_dict(),  # 保存模型的状态字典
    'optimizer_state_dict': optimizer.state_dict(),  # 保存优化器的状态字典
    'loss': train_loss,                     # 保存当前训练的损失值
}, 'model_checkpoint.tar')