基于Transformer的机器翻译模型

目录

一、实验原理

二、实验步骤

三、实验内容

1、相关库的导入

2、构建数据集

3、准备分词器

4、构建TorchText词汇对象并将句子转换为Torch张量

5、创建在训练过程中迭代的数据加载器（DataLoader）对象

6、序列到序列的Transformer模型

7、模型训练

8、尝试使用训练好的模型翻译一个日语句子

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

四、实验总结

五、实验代码

---

一、实验原理

Transformer结构

Transformer模型是一种使用注意力机制来进行序列建模的架构，最初被提出用于自然语言处理任务，特别是机器翻译。

Transformer模型的核心思想是完全基于注意力机制，而不依赖于传统的循环神经网络（RNN）或卷积神经网络（CNN）。它由编码器和解码器两部分组成，每部分都由多层堆叠的注意力模块和前馈神经网络模块构成。

二、实验步骤

1. 导入所需要的库：我们需要安装torch等代码中所需要的库，并且保证版本的适配性。
2. 数据集的构建：我们下载导入网络上的日英数据集。
3. 准备分词器：使用分词器（如SentencePiece）构建日语和英语的词汇表。
4. 模型定义：定义Transformer模型，包括编码器和解码器，以及嵌入层、注意力机制和前馈网络。
5. 模型训练：设置训练参数，如批处理大小、学习率等。使用训练数据训练Transformer模型。监控训练过程中的损失和性能指标。
6. 模型评估：使用验证集评估模型的翻译能力，如计算BLEU等指标。
7. 模型测试：我们找个一个简单的例子输入，测试翻译的效果
8. 模型保存：使用Pickle保存词汇表对象（en_vocab和ja_vocab）。使用PyTorch的保存函数保存模型，以便后续推断或恢复训练使用。

三、实验内容

1、相关库的导入

在实验开始前，我们需要保证自己所有的库都已经安装并且版本匹配。像我就遇到了版本库版本的问题。torchtext默认下载的版本是0.18.0，但是会报错，因为代码中很多功能更新了，所以需要降级到0.4.0，也是最常用的版本。

代码如下：

import math  # 导入数学模块，提供基本的数学运算函数
import torchtext  # 导入torchtext库，用于处理文本数据
import torch  # 导入PyTorch库，用于深度学习模型的构建和训练
import torch.nn as nn  # 导入PyTorch的神经网络模块
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  # 导入Counter类，用于计数
from torchtext.vocab import Vocab  # 导入Vocab类，用于词汇表管理
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer# 导入Transformer相关的编码器和解码器模块
import io  # 导入io模块，用于文件操作
import time  # 导入time模块，用于计时
import pandas as pd  # 导入pandas库，用于数据处理
import numpy as np  # 导入numpy库，用于科学计算
import pickle  # 导入pickle模块，用于序列化和反序列化对象
import tqdm  # 导入tqdm库，用于显示进度条
import sentencepiece as spm  # 导入sentencepiece库，用于子词分割

# 设置随机种子，以确保结果的可复现性
torch.manual_seed(0)
# 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(torch.cuda.get_device_name(0)) # 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

2、构建数据集

在本文中，我们将使用从JParaCrawl下载的日英平行数据集。http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl。http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl%5D%E3%80%82该数据集被描述为“由NTT创建的最大公开可用的英语-日语平行语料库”。它主要通过爬取网页并自动对齐平行句子来创建。

代码如下：

# 读取名为'zh-ja.bicleaner05.txt'的CSV文件，使用制表符作为字段分隔符进行分割
# engine='python'表示使用Python引擎来解析CSV文件，header=None表示该文件没有列名
df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
# 从DataFrame中提取第2列的数据，并将其转换为列表形式
trainen = df[2].values.tolist()#[:10000]
# 从DataFrame中提取第3列的数据，并将其转换为列表形式
trainja = df[3].values.tolist()#[:10000]
# 下面两行代码是注释掉的，如果取消注释，则会删除trainen和trainja列表中索引为5972的元素
# trainen.pop(5972)
# trainja.pop(5972)

在导入所有日语句子及其对应的英语句子后，我删除了数据集中最后一条因为它有缺失值。总的来说，训练集中的句子数量在trainen和trainja中都是5,973,071。然而，出于学习目的，通常建议对数据进行采样，并确保所有操作按预期工作，然后再一次性使用全部数据，以节省时间。

代码如下：

# 打印trainen列表中索引为500的元素
print(trainen[500])
# 打印trainja列表中索引为500的元素
print(trainja[500])

结果输出：

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

3、准备分词器

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

代码如下：

# 加载英文的SentencePiece模型，用于对英文文本进行分词
# model_file 参数指定了模型文件路径，这里是 'spm.en.nopretok.model'
en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')

# 加载日文的SentencePiece模型，用于对日文文本进行分词
# model_file 参数指定了模型文件路径，这里是 'spm.ja.nopretok.model'
ja_tokenizer = spm.SentencePieceProcessor(model_file='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')

结果输出：

4、构建TorchText词汇对象并将句子转换为Torch张量

利用分词器和原始句子，我们可以构建来自TorchText的词汇（Vocab）对象。这个过程可能需要几秒钟或几分钟，具体取决于数据集的大小和计算能力。不同的分词器也会影响构建词汇所需的时间。我尝试了几种其他日语分词器，但是发现SentencePiece 对我来说工作得很好，并且速度足够快。
代码如下：

def build_vocab(sentences, tokenizer):
    # 创建一个空的计数器对象
    counter = Counter()
    # 遍历每个句子，对分词后的结果进行计数
    for sentence in sentences:
        counter.update(tokenizer.encode(sentence, out_type=str))
    # 使用Counter对象创建一个新的词汇表，并指定特殊标记
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
# 使用 build_vocab 函数分别构建日文和英文的词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer) # trainja 是包含日文句子的列表
en_vocab = build_vocab(trainen, en_tokenizer) # trainen 是包含英文句子的列表

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

代码如下：

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)

5、创建在训练过程中迭代的数据加载器（DataLoader）对象

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

代码如下：

# 定义批处理大小
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:
        # 在日文句子的开头添加起始符，在末尾添加结束符，并将其放入日文句子张量列表中
        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))
    # 对日文句子张量列表进行填充，使用 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)

6、序列到序列的Transformer模型

接下来的几段代码和文字说明（用斜体书写）摘自原始的PyTorch教程 https://pytorch.org/tutorials/beginner/translation_transformer.html]。https://pytorch.org/tutorials/beginner/translation_transformer.html%5D%E3%80%82

除了将BATCH_SIZE和单词de_vocab改为ja_vocab外，我没有做任何更改。

Transformer是一种序列到序列（Seq2Seq）模型，介绍于论文“Attention is all you need”，用于解决机器翻译任务。Transformer模型包括一个编码器和一个解码器块，每个块包含固定数量的层。

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

代码如下：

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):
        """
        初始化 Seq2Seq Transformer 模型

        参数:
        num_encoder_layers (int): 编码器层数
        num_decoder_layers (int): 解码器层数
        emb_size (int): 嵌入向量大小
        src_vocab_size (int): 源语言词汇表大小
        tgt_vocab_size (int): 目标语言词汇表大小
        dim_feedforward (int, optional): 前馈神经网络的隐藏层大小，默认值为 512
        dropout (float, optional): dropout 概率，默认值为 0.1
        """
        super(Seq2SeqTransformer, self).__init__()
        # 定义编码器层
        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)
        
        # 定义解码器层
        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 (Tensor): 源语言输入序列
        trg (Tensor): 目标语言输入序列
        src_mask (Tensor): 源语言序列的掩码
        tgt_mask (Tensor): 目标语言序列的掩码
        src_padding_mask (Tensor): 源语言序列的填充值掩码
        tgt_padding_mask (Tensor): 目标语言序列的填充值掩码
        memory_key_padding_mask (Tensor): 编码器输出的填充值掩码

        返回:
        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):
        """
        编码函数

        参数:
        src (Tensor): 源语言输入序列
        src_mask (Tensor): 源语言序列的掩码

        返回:
        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):
        """
        解码函数

        参数:
        tgt (Tensor): 目标语言输入序列
        memory (Tensor): 编码器的输出
        tgt_mask (Tensor): 目标语言序列的掩码

        返回:
        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):
        """
        初始化位置编码模块。
        
        参数:
        emb_size (int): 嵌入向量的维度。
        dropout (float): dropout 概率，用于防止过拟合。
        maxlen (int): 序列的最大长度，默认值为 5000。
        """
        # 计算位置编码中的分母部分
        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)
        
        # 定义dropout层
        self.dropout = nn.Dropout(dropout)
        
        # 注册位置嵌入矩阵为buffer，使其在模型保存和加载时不会作为参数
        self.register_buffer('pos_embedding', pos_embedding)


    def forward(self, token_embedding: Tensor):
        """
        前向传播函数，将位置编码添加到token嵌入上，并应用dropout。
        
        参数:
        token_embedding (Tensor): 输入的token嵌入张量，形状为 [seq_len, batch_size, emb_size]。
        
        返回:
        Tensor: 叠加位置编码后的token嵌入张量。
        """
        # 将位置编码添加到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):
        """
        初始化token嵌入模块。
        
        参数:
        vocab_size (int): 词汇表的大小。
        emb_size (int): 嵌入向量的维度。
        """
        super(TokenEmbedding, self).__init__()
        # 定义嵌入层
        self.embedding = nn.Embedding(vocab_size, emb_size)
        # 存储嵌入向量的维度
        self.emb_size = emb_size
    def forward(self, tokens: Tensor):
        """
        前向传播函数，获取tokens的嵌入表示并进行缩放。
        
        参数:
        tokens (Tensor): 输入的token张量，形状通常为 [seq_len, batch_size]。
        
        返回:
        Tensor: 缩放后的token嵌入张量，形状为 [seq_len, batch_size, emb_size]。
        """
        # 获取token的嵌入表示，并乘以嵌入维度的平方根进行缩放
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建一个后续词掩码，以阻止目标词与其后续词产生关联。我们还创建了用于屏蔽源和目标填充掩码。

代码如下：

# 定义生成方形后续掩码的函数
def generate_square_subsequent_mask(sz):
    """
    生成一个大小为 sz x sz 的方形掩码矩阵，其中上三角部分为 0，其他部分为 -inf。
    这个掩码用于在解码过程中屏蔽未来的位置，以保证自回归模型只能看到当前和之前的 token。

    参数:
    sz (int): 掩码矩阵的大小 (边长)

    返回:
    torch.Tensor: 一个大小为 sz x sz 的掩码矩阵
    """
    # 创建一个上三角矩阵，值为 1，表示允许的位置；之后转置矩阵
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 把上三角部分保留为 0，其他部分填充为 -inf，表示不允许的位置
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

# 定义创建掩码的函数
def create_mask(src, tgt):
    """
    为源序列和目标序列创建必要的掩码，包括源掩码、目标掩码、源填充掩码和目标填充掩码。

    参数:
    src (torch.Tensor): 源序列张量，形状为 (src_seq_len, batch_size)
    tgt (torch.Tensor): 目标序列张量，形状为 (tgt_seq_len, batch_size)

    返回:
    tuple: 包含以下四个元素的元组：
        - src_mask (torch.Tensor): 源序列掩码矩阵
        - tgt_mask (torch.Tensor): 目标序列掩码矩阵
        - src_padding_mask (torch.Tensor): 源序列填充掩码矩阵
        - tgt_padding_mask (torch.Tensor): 目标序列填充掩码矩阵
    """
    # 获取源序列和目标序列的长度
    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]

    # 生成目标序列的掩码，防止解码时看到未来的 token
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)

    # 生成源序列的掩码，这里全为 False，因为不需要屏蔽任何位置
    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

因为平台运行不了，接下来的模型是在自己的电脑上运行的了。我将NUM_ENCODER_LAYERS 和 NUM_DECODER_LAYERS 设置为3，NHEAD设置8，EMB_SIZE设置为512。

代码如下：

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)

# 将模型移动到指定的设备上
transformer = transformer.to(device)

# 定义交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# 定义优化器
optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)
# 定义训练函数
def train_epoch(model, train_iter, optimizer):
    """
    对模型进行一个 epoch 的训练，并返回平均损失。

    参数:
    model: 待训练的模型
    train_iter: 训练数据迭代器
    optimizer: 优化器

    返回:
    float: 平均损失
    """
    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: 待评估的模型
    val_iter: 验证数据迭代器

    返回:
    float: 平均损失
    """
    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)

7、模型训练

当我使用JParaCrawl的完整句子列表训练模型时，每种语言大约有590万句子，使用VSCode远程连接服务器，GPU类型为NVIDIA GeForce RTX 3090 GPU，每个epoch大约需要5分钟左右。

最后，在准备好必要的类和函数之后，我们就可以开始训练我们的模型了。不用说，完成训练所需的时间可能会因多种因素而有很大差异，例如计算能力、参数和数据集的大小。

代码如下：

# tqdm.tqdm用于生成进度条
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
    # 记录一个epoch开始的时间
    start_time = time.time()
    # 调用训练函数进行模型训练，并返回该epoch的平均训练损失
    train_loss = train_epoch(transformer, train_iter, optimizer)
    # 记录一个epoch结束的时间
    end_time = time.time()
    # 打印当前epoch的编号、训练损失以及耗时
    print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
            f"Epoch time = {(end_time - start_time):.3f}s"))

结果输出：

从结果可以看出，随着批次的增加，损失值在不断减小，但在第十个epoch后下降的速度就变得很平缓了。

8、尝试使用训练好的模型翻译一个日语句子

在3090显卡上训练完模型后，为了保证输出的便捷和美观，我将训练好的模型导入到pycharm里进行后面的翻译工作。

首先，我们创建用于翻译新句子的函数，包括以下步骤：获取日语句子、进行标记化处理、转换为张量、推理，然后将结果解码回英语句子。

代码如下：

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)
        # 创建目标语言的mask
        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()
    # 对源语言进行tokenize，并在开头和结尾加上特殊符号
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)]+ [EOS_IDX]
    num_tokens = len(tokens)
    # 将源语言转换为张量，并创建对应的mask
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1) )
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
    # 使用greedy_decode函数生成目标语言的翻译结果
    tgt_tokens = greedy_decode(model,  src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
    # 将生成的目标语言token转换为文本并进行后处理
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

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

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

结果输出：

我自己搜索，这句话的大意是"HS Code 8515: Equipment for soldering, brazing, or welding (电气式（包括电加热气体式）)."。

除了给定的例子，我还自己测试了几句话：

这两句话的中文分别是：

• 董老师很帅
• 人工智能很厉害

翻译出来的是

由此可见，翻译的并不是很好。

trainen.pop(5)

trainja.pop(5)

结果输出：

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

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

代码如下：

import pickle
# 打开一个文件，用于存储数据
file = open('en_vocab.pkl', 'wb')
# 将数据(en_vocab)写入文件中
pickle.dump(en_vocab, file)
# 关闭文件
file.close()
# 打开一个文件，用于存储数据
file = open('ja_vocab.pkl', 'wb')
# 将数据(ja_vocab)写入文件中
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')  # 保存文件的路径和名称

四、实验总结

在本文中，我们介绍了Transformer模型，对Transformer模型的结构有了更深刻的认识和理解。通过对Transformer在机器翻译中的运用，我们可以更好的了解到自然语言处理中的核心技术。对于此次内容，我们需要掌握其构建的流程。

核心内容：

• Transformer的结构
• 构建TorchText词汇对象
• Transformer模型构建

问题总结：本次实验对于电脑的要求较高，运行时间较长；所需要的库需要注意版本匹配问题；原代码中存在代码缩进问题，需要逐行检查修改。

通过实现和训练基于Transformer的机器翻译模型，我们展示了如何在实际应用中进行机器翻译任务。本次实验不仅可以帮助我们更好的理解了Transformer等核心模型的应用，更让我们更好的掌握自然语言处理的核心技术。

五、实验代码

import math  # 导入数学模块，提供基本的数学运算函数
import torchtext  # 导入torchtext库，用于处理文本数据
import torch  # 导入PyTorch库，用于深度学习模型的构建和训练
import torch.nn as nn  # 导入PyTorch的神经网络模块
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  # 导入Counter类，用于计数
from torchtext.vocab import Vocab  # 导入Vocab类，用于词汇表管理
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer# 导入Transformer相关的编码器和解码器模块
import io  # 导入io模块，用于文件操作
import time  # 导入time模块，用于计时
import pandas as pd  # 导入pandas库，用于数据处理
import numpy as np  # 导入numpy库，用于科学计算
import pickle  # 导入pickle模块，用于序列化和反序列化对象
import tqdm  # 导入tqdm库，用于显示进度条
import sentencepiece as spm  # 导入sentencepiece库，用于子词分割

# 设置随机种子，以确保结果的可复现性
torch.manual_seed(0)
# 检查是否有可用的GPU，如果有则使用GPU，否则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#print(torch.cuda.get_device_name(0)) # 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

device

# 读取名为'zh-ja.bicleaner05.txt'的CSV文件，使用制表符作为字段分隔符进行分割
# engine='python'表示使用Python引擎来解析CSV文件，header=None表示该文件没有列名
df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
# 从DataFrame中提取第2列的数据，并将其转换为列表形式
trainen = df[2].values.tolist()#[:10000]
# 从DataFrame中提取第3列的数据，并将其转换为列表形式
trainja = df[3].values.tolist()#[:10000]
# 下面两行代码是注释掉的，如果取消注释，则会删除trainen和trainja列表中索引为5972的元素
# trainen.pop(5972)
# trainja.pop(5972)

# 打印trainen列表中索引为500的元素
print(trainen[500])
# 打印trainja列表中索引为500的元素
print(trainja[500])

# 加载英文的SentencePiece模型，用于对英文文本进行分词
# model_file 参数指定了模型文件路径，这里是 'spm.en.nopretok.model'
en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')

# 加载日文的SentencePiece模型，用于对日文文本进行分词
# model_file 参数指定了模型文件路径，这里是 'spm.ja.nopretok.model'
ja_tokenizer = spm.SentencePieceProcessor(model_file='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')

def build_vocab(sentences, tokenizer):
    # 创建一个空的计数器对象
    counter = Counter()
    # 遍历每个句子，对分词后的结果进行计数
    for sentence in sentences:
        counter.update(tokenizer.encode(sentence, out_type=str))
    # 使用Counter对象创建一个新的词汇表，并指定特殊标记
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
# 使用 build_vocab 函数分别构建日文和英文的词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer) # trainja 是包含日文句子的列表
en_vocab = build_vocab(trainen, en_tokenizer) # trainen 是包含英文句子的列表

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)

# 定义批处理大小
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:
        # 在日文句子的开头添加起始符，在末尾添加结束符，并将其放入日文句子张量列表中
        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))
    # 对日文句子张量列表进行填充，使用 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)

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):
        """
        初始化 Seq2Seq Transformer 模型

        参数:
        num_encoder_layers (int): 编码器层数
        num_decoder_layers (int): 解码器层数
        emb_size (int): 嵌入向量大小
        src_vocab_size (int): 源语言词汇表大小
        tgt_vocab_size (int): 目标语言词汇表大小
        dim_feedforward (int, optional): 前馈神经网络的隐藏层大小，默认值为 512
        dropout (float, optional): dropout 概率，默认值为 0.1
        """
        super(Seq2SeqTransformer, self).__init__()
        # 定义编码器层
        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)
        
        # 定义解码器层
        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 (Tensor): 源语言输入序列
        trg (Tensor): 目标语言输入序列
        src_mask (Tensor): 源语言序列的掩码
        tgt_mask (Tensor): 目标语言序列的掩码
        src_padding_mask (Tensor): 源语言序列的填充值掩码
        tgt_padding_mask (Tensor): 目标语言序列的填充值掩码
        memory_key_padding_mask (Tensor): 编码器输出的填充值掩码

        返回:
        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):
        """
        编码函数

        参数:
        src (Tensor): 源语言输入序列
        src_mask (Tensor): 源语言序列的掩码

        返回:
        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):
        """
        解码函数

        参数:
        tgt (Tensor): 目标语言输入序列
        memory (Tensor): 编码器的输出
        tgt_mask (Tensor): 目标语言序列的掩码

        返回:
        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):
        """
        初始化位置编码模块。
        
        参数:
        emb_size (int): 嵌入向量的维度。
        dropout (float): dropout 概率，用于防止过拟合。
        maxlen (int): 序列的最大长度，默认值为 5000。
        """
        # 计算位置编码中的分母部分
        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)
        
        # 定义dropout层
        self.dropout = nn.Dropout(dropout)
        
        # 注册位置嵌入矩阵为buffer，使其在模型保存和加载时不会作为参数
        self.register_buffer('pos_embedding', pos_embedding)


    def forward(self, token_embedding: Tensor):
        """
        前向传播函数，将位置编码添加到token嵌入上，并应用dropout。
        
        参数:
        token_embedding (Tensor): 输入的token嵌入张量，形状为 [seq_len, batch_size, emb_size]。
        
        返回:
        Tensor: 叠加位置编码后的token嵌入张量。
        """
        # 将位置编码添加到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):
        """
        初始化token嵌入模块。
        
        参数:
        vocab_size (int): 词汇表的大小。
        emb_size (int): 嵌入向量的维度。
        """
        super(TokenEmbedding, self).__init__()
        # 定义嵌入层
        self.embedding = nn.Embedding(vocab_size, emb_size)
        # 存储嵌入向量的维度
        self.emb_size = emb_size
    def forward(self, tokens: Tensor):
        """
        前向传播函数，获取tokens的嵌入表示并进行缩放。
        
        参数:
        tokens (Tensor): 输入的token张量，形状通常为 [seq_len, batch_size]。
        
        返回:
        Tensor: 缩放后的token嵌入张量，形状为 [seq_len, batch_size, emb_size]。
        """
        # 获取token的嵌入表示，并乘以嵌入维度的平方根进行缩放
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

# 定义生成方形后续掩码的函数
def generate_square_subsequent_mask(sz):
    """
    生成一个大小为 sz x sz 的方形掩码矩阵，其中上三角部分为 0，其他部分为 -inf。
    这个掩码用于在解码过程中屏蔽未来的位置，以保证自回归模型只能看到当前和之前的 token。

    参数:
    sz (int): 掩码矩阵的大小 (边长)

    返回:
    torch.Tensor: 一个大小为 sz x sz 的掩码矩阵
    """
    # 创建一个上三角矩阵，值为 1，表示允许的位置；之后转置矩阵
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 把上三角部分保留为 0，其他部分填充为 -inf，表示不允许的位置
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

# 定义创建掩码的函数
def create_mask(src, tgt):
    """
    为源序列和目标序列创建必要的掩码，包括源掩码、目标掩码、源填充掩码和目标填充掩码。

    参数:
    src (torch.Tensor): 源序列张量，形状为 (src_seq_len, batch_size)
    tgt (torch.Tensor): 目标序列张量，形状为 (tgt_seq_len, batch_size)

    返回:
    tuple: 包含以下四个元素的元组：
        - src_mask (torch.Tensor): 源序列掩码矩阵
        - tgt_mask (torch.Tensor): 目标序列掩码矩阵
        - src_padding_mask (torch.Tensor): 源序列填充掩码矩阵
        - tgt_padding_mask (torch.Tensor): 目标序列填充掩码矩阵
    """
    # 获取源序列和目标序列的长度
    src_seq_len = src.shape[0]
    tgt_seq_len = tgt.shape[0]

    # 生成目标序列的掩码，防止解码时看到未来的 token
    tgt_mask = generate_square_subsequent_mask(tgt_seq_len)

    # 生成源序列的掩码，这里全为 False，因为不需要屏蔽任何位置
    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)

# 将模型移动到指定的设备上
transformer = transformer.to(device)

# 定义交叉熵损失函数
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# 定义优化器
optimizer = torch.optim.Adam(
    transformer.parameters(), lr=0.0001, betas=(0.9, 0.98), eps=1e-9
)
# 定义训练函数
def train_epoch(model, train_iter, optimizer):
    """
    对模型进行一个 epoch 的训练，并返回平均损失。

    参数:
    model: 待训练的模型
    train_iter: 训练数据迭代器
    optimizer: 优化器

    返回:
    float: 平均损失
    """
    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: 待评估的模型
    val_iter: 验证数据迭代器

    返回:
    float: 平均损失
    """
    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)

# tqdm.tqdm用于生成进度条
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
    # 记录一个epoch开始的时间
    start_time = time.time()
    # 调用训练函数进行模型训练，并返回该epoch的平均训练损失
    train_loss = train_epoch(transformer, train_iter, optimizer)
    # 记录一个epoch结束的时间
    end_time = time.time()
    # 打印当前epoch的编号、训练损失以及耗时
    print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
            f"Epoch time = {(end_time - start_time):.3f}s"))

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)
        # 创建目标语言的mask
        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()
    # 对源语言进行tokenize，并在开头和结尾加上特殊符号
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)]+ [EOS_IDX]
    num_tokens = len(tokens)
    # 将源语言转换为张量，并创建对应的mask
    src = (torch.LongTensor(tokens).reshape(num_tokens, 1) )
    src_mask = (torch.zeros(num_tokens, num_tokens)).type(torch.bool)
    # 使用greedy_decode函数生成目标语言的翻译结果
    tgt_tokens = greedy_decode(model,  src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
    # 将生成的目标语言token转换为文本并进行后处理
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

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

trainen.pop(5)
trainja.pop(5)

import pickle
# 打开一个文件，用于存储数据
file = open('en_vocab.pkl', 'wb')
# 将数据(en_vocab)写入文件中
pickle.dump(en_vocab, file)
# 关闭文件
file.close()
# 打开一个文件，用于存储数据
file = open('ja_vocab.pkl', 'wb')
# 将数据(ja_vocab)写入文件中
pickle.dump(ja_vocab, file)
# 关闭文件
file.close()

# save model for inference
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')  # 保存文件的路径和名称