基于Transformer和PyTorch的日语——英语机器翻译模型

一：导入所需要的包

首先，请确保我们的系统中安装了下面的包。如果发现有些包没有安装，一定要安装它们。

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)) ## 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

device

二：获取平行数据集

在本教程中，我们将使用从JParaCrawl下载的日-英平行数据集。http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl该数据集被描述为“NTT创建的最大的公开可用的英-日平行语料库。它主要是通过大量爬取网络并自动对齐平行句子来创建的。”您也可以在这里查看论文。

df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)  # 读取文件zh-ja.bicleaner05.txt，使用Tab作为分隔符，无表头
trainen = df[2].values.tolist()  # 获取数据框的第三列并转换为列表
trainja = df[3].values.tolist()  # 获取数据框的第四列并转换为列表
# trainen.pop(5972)  # 移除trainen列表中的第5972个元素（如果需要）
# trainja.pop(5972)  # 移除trainja列表中的第5972个元素（如果需要）

在导入所有日语及其英语对应内容后，我删除了数据集中的最后一条数据，因为它存在缺失值。总共，trainen和trainja中的句子数量为5,973,071个。然而，出于学习目的，通常建议先对数据进行抽样，确保一切按预期运行，以节省时间，然后再一次性使用所有数据。

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

print(trainen[500])
print(trainja[500])

我们还可以采用其他不同的平行数据集来配合本文的学习，只需确保我们可以将数据加工成上面所示的两个字符串列表，包含日语和英语句子。

三：准备tokenizers

与英语或其他字母语言不同，日语句子中没有空格来分隔单词。我们可以使用JParaCrawl提供的标记器，该工具是使用SentencePiece为日语和英语创建的。

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

加载tokenizers后，你可以通过执行以下代码来测试它们。

en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type='str')# 使用英文的SentencePiece分词器对句子进行编码，将其分割成子词（子词作为字符串类型输出）

ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。", out_type='str')# 使用日文的SentencePiece分词器对句子进行编码，将其分割成子词（子词作为字符串类型输出）

四：构建TorchText词汇对象并把句子转换成Torch张量

使用标记器和原始句子，我们接下来构建从TorchText中导入的Vocab对象。这个过程可能需要几秒到几分钟的时间，具体取决于数据集的大小和计算能力。不同的标记器也会影响构建词汇表所需的时间，我尝试了几种其他的日语标记器，但SentencePiece似乎对我来说工作得很好而且足够快。

def build_vocab(sentences, tokenizer):
    counter = 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 = []  # 创建一个空列表，用于存储处理后的数据
    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设置为16以防止“cuda内存不足”，但这取决于各种因素，例如您的机器内存容量、数据大小等，因此请根据您的需求自由更改批处理大小（注意：使用Multi30k德英数据集的PyTorch教程将批处理大小设置为128。）

BATCH_SIZE = 8  # 设置批处理大小为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))
    # 对日文和英文句子进行填充，使它们具有相同的长度
    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  # 返回填充后的日文和英文批处理数据

# 创建数据加载器，使用训练数据，每个批次包含BATCH_SIZE个样本，打乱数据，并使用generate_batch进行数据处理
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

六：Sequence-to-sequence Transformer

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

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

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

from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)  # 导入Transformer相关模块

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)  # 创建Transformer编码器
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)  # 创建解码器层
        self.transformer_decoder = TransformerDecoder(decoder_layer, num_layers=num_decoder_layers)  # 创建Transformer解码器

        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)  # 解码目标语言输入

Text tokens通过使用token嵌入来表示。将位置编码添加到token嵌入中，以引入单词顺序的概念。

import math
import torch
import torch.nn as nn
from torch import Tensor

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):
        # 在输入的词嵌入上加上位置编码，并进行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):
        # 将输入的词索引转换为词嵌入表示，并乘以sqrt(词嵌入维度)
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建了一个后续词掩码，以防止目标词关注其后续词。我们还创建了用于掩盖源和目标填充token的掩码。

def generate_square_subsequent_mask(sz):
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)  # 创建一个上三角矩阵
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))  # 将上三角矩阵中的0替换为负无穷，1替换为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)  # 生成目标序列的自注意力掩码
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)  # 创建源序列的掩码矩阵（全为False）

    src_padding_mask = (src == PAD_IDX).transpose(0, 1)  # 创建源序列的填充掩码，PAD_IDX是填充标记的索引
    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)  # 创建目标序列的填充掩码，PAD_IDX是填充标记的索引
    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  # 训练轮数

# 初始化Transformer模型
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)  # 将模型移动到GPU（如果可用）

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
)  # 定义优化器（Adam优化器）

def train_epoch(model, train_iter, optimizer):
    model.train()  # 设置模型为训练模式
    losses = 0
    for idx, (src, tgt) in enumerate(train_iter):
        src = src.to(device)  # 将源语言数据移动到GPU（如果可用）
        tgt = tgt.to(device)  # 将目标语言数据移动到GPU（如果可用）

        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)  # 将源语言数据移动到GPU（如果可用）
        tgt = tgt.to(device)  # 将目标语言数据移动到GPU（如果可用）

        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)  # 返回平均损失

七：开始训练

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

import tqdm
import time

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS + 1)):
    start_time = time.time()  # 记录当前epoch开始时间
    train_loss = train_epoch(transformer, train_iter, optimizer)  # 训练模型并计算损失
    end_time = time.time()  # 记录当前epoch结束时间
    # 打印当前epoch的训练损失和epoch运行时间
    print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, 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)
        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()  # 使用贪婪解码生成目标语言序列
    translation = " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")  # 将生成的目标语言序列转换为文本
    return translation

然后，我们只需调用translate函数并传递所需的参数。

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

trainen.pop(5)

trainja.pop(5)

九：保存词汇表对象和训练过的模型

最后，在训练完成后，我们首先使用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的保存和加载功能将模型保存以供以后使用。通常，根据我们以后想要如何使用它们，有两种方式可以保存模型。第一种是仅用于推理，我们可以稍后加载模型并使用它将日语翻译成英语。

# save model for inference
torch.save(transformer.state_dict(), 'inference_model')

第二种也是用于推理，但也适用于我们稍后想要加载模型并希望恢复训练的情况。

# save model + checkpoint to resume training later
torch.save({
    'epoch': NUM_EPOCHS,  # 保存当前训练的轮数
    'model_state_dict': transformer.state_dict(),  # 保存模型的状态字典
    'optimizer_state_dict': optimizer.state_dict(),  # 保存优化器的状态字典
    'loss': train_loss,  # 保存当前训练的损失值
}, 'model_checkpoint.tar')  # 文件名为'model_checkpoint.tar'