使用Transformer和PyTorch进行日文-中文机器翻译模型

本文将讲解如何使用Jupyter Notebook、PyTorch、Torchtext和SentencePiece来构建日文-中文机器翻译模型。我们将使用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')
# 如果你有GPU，请在你自己的电脑上尝试运行这一套代码
# print(torch.cuda.get_device_name(0))
device

获取平行数据集

在这个教程中，我们将使用从JParaCrawl下载的日文-英语平行数据集。这是由NTT创建的目前最大的公开的英日平行语料库。它主要通过爬取网页并自动对齐平行句子来创建。你可以在这里看到相关论文。

df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainen = df[2].values.tolist() # [:10000]
trainja = df[3].values.tolist() # [:10000]
# trainen.pop(5972)
# trainja.pop(5972)

导入所有日文和对应的英文句子后，我删除了数据集中最后一个数据，因为它有一个缺失值。总的来说，trainen和trainja中句子的数量是5,973,071。然而，为了学习目的，通常建议对数据进行采样，并确保一切正常工作后，再一次性使用所有数据，以节省时间。

数据集中包含的句子示例如下：

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 ...

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

准备分词器

在日文中，不像英文等字母语言，句子中没有空格来分隔单词。因此，我们使用由SentencePiece提供的JParaCrawl分词器来对日文和英文进行分词。你可以访问JParaCrawl网站下载它们的模型。

en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.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')

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

使用分词器和原始句子，我们构建从TorchText导入的词汇对象。这个过程可能需要几秒钟到几分钟，具体取决于数据集的大小和计算能力。不同的分词器也会影响构建词汇表所需的时间。尽管尝试了几种其他的日文分词器，但SentencePiece在速度和效果上都表现良好。

def build_vocab(sentences, tokenizer):
    """
    构建词汇表的函数。

    参数：
    - sentences: 一个包含多个句子的列表。
    - tokenizer: 用于将句子分词的分词器。

    返回：
    - 一个包含词汇和频率的 Vocab 对象。
    """
    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):
    """
    处理日语和英语句子列表，将其转换为张量。

    参数：
    - ja: 包含日语句子的列表。
    - en: 包含英语句子的列表。

    返回：
    - data: 包含 (ja_tensor_, en_tensor_) 元组的列表，其中每个元组对应一对日语和英语句子的张量表示。
    """
    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对象进行训练

在训练过程中，我们需要构建一个DataLoader对象来迭代数据。在这里，我们将批处理大小（BATCH_SIZE）设置为16，以防止“cuda out of memory”错误。不过，这取决于你的机器内存容量、数据大小等因素，所以你可以根据需要随意更改批处理大小（注：PyTorch的教程中使用Multi30k德英数据集时，批处理大小为128）。

BATCH_SIZE = 8
# 获取特殊 token 的索引
PAD_IDX = ja_vocab['<pad>']
BOS_IDX = ja_vocab['<bos>']
EOS_IDX = ja_vocab['<eos>']

def generate_batch(data_batch):
    """
    生成一个批处理数据的函数。

    参数：
    - data_batch: 包含 (ja_item, en_item) 元组的列表，每个元组是日语和英语句子的张量表示。

    返回：
    - ja_batch: 填充后的日语句子张量。
    - en_batch: 填充后的英语句子张量。
    """
    ja_batch, en_batch = []  # 初始化日语和英语批处理列表
    
    # 遍历批处理数据中的每个句子对
    for (ja_item, en_item) in data_batch:
        # 在日语句子前后添加起始和结束 token，并添加到批处理列表中
        ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
        # 在英语句子前后添加起始和结束 token，并添加到批处理列表中
        en_batch.append(torch.cat([torch.tensor([BOS_IDX]), en_item, torch.tensor([EOS_IDX])], dim=0))
    
    # 使用 pad_sequence 对批处理列表进行填充
    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)

构建Sequence-to-Sequence Transformer模型

接下来的代码和文字说明（以斜体书写的部分）取自原始的PyTorch教程[https://pytorch.org/tutorials/beginner/translation_transformer.html]。除了BATCH_SIZE和词汇对象从de_vocab改为ja_vocab外，我没有做任何更改。

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

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

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: 源语言输入张量
        - trg: 目标语言输入张量
        - src_mask: 源语言掩码
        - tgt_mask: 目标语言掩码
        - src_padding_mask: 源语言填充掩码
        - tgt_padding_mask: 目标语言填充掩码
        - memory_key_padding_mask: 编码器输出的填充掩码

        返回：
        - 目标语言的生成结果
        """
        # 获取源语言和目标语言的嵌入表示，并加上位置编码
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
    

位置编码和词嵌入

文本token使用词嵌入表示。位置编码加到词嵌入中引入单词顺序的概念。

class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout: float, maxlen: int = 5000):
        """
        初始化位置编码模块。

        参数：
        - emb_size: 词嵌入的维度。
        - dropout: 位置编码后的 dropout 概率。
        - maxlen: 最大位置的长度，默认值为 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)

        # 定义 dropout 层
        self.dropout = nn.Dropout(dropout)
        
        # 注册位置编码矩阵为缓冲区
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        """
        前向传播函数，给输入的 token embedding 加上位置编码。

        参数：
        - token_embedding: 输入的词嵌入张量。

        返回：
        - 加上位置编码并应用

创建掩码

我们创建一个后续单词掩码，以阻止目标单词关注其后续单词。我们还为源和目标填充token创建掩码。

def generate_square_subsequent_mask(sz):
    """
    生成一个大小为 sz 的下三角掩码张量，用于 Transformer 解码器中的自注意力机制。

    参数：
    - sz: 掩码的大小，即序列的长度。

    返回：
    - 一个大小为 (sz, sz) 的掩码张量，上三角部分填充为 -inf，下三角部分填充为 0.0。
    """
    # 创建一个上三角为1，下三角为0的张量
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将上三角部分填充为 -inf，表示这些位置在注意力机制中不可见
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    return mask

def create_mask(src, tgt):
    """
    为源序列和目标序列创建掩码，用于 Transformer 模型中的注意力机制。

    参数：
    - src: 源序列张量，形状为 (src_seq_len, batch_size)。
    - tgt: 目标序列张量，形状为 (tgt_seq_len, batch_size)。

    返回：
    - src_mask: 源序列的掩码，形状为 (src_seq_len, src_seq_len)。
    - tgt_mask: 目标序列的掩码，形状为 (tgt_seq_len, tgt_seq_len)。
    - src_padding_mask: 源序列的填充掩码，形状为 (batch_size, src_seq_len)。
    - tgt_padding_mask: 目标序列的填充掩码，形状为 (batch_size, tgt_seq_len)。
    """
    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)

    # 生成源序列和目标序列的填充掩码
    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

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

这里由于服务器计算能力有限，按照以下配置可以进行训练，但效果可能不佳。如果想看到训练效果，请使用你自己的带GPU的电脑运行这套代码。

当使用自己的GPU时，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)

# 初始化模型参数
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)

# 定义优化器
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: 优化器。

    返回：
    - 平均损失。
    """
    model.train()  # 设置模型为训练模式
    losses = 0  # 初始化损失

    # 遍历训练数据
    for idx, (src, tgt) in enumerate(train_iter):
        src = src.to(device)  # 将源序列移动到设备
        tgt = tgt.to(device)  # 将目标序列移动到设备

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

        # 创建掩码
        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:, :]  # 目标输出（去掉第一个 token）
        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: 验证数据的迭代器。

    返回：
    - 平均损失。
    """
    model.eval()  # 设置模型为评估模式
    losses = 0  # 初始化损失

    # 遍历验证数据
    for idx, (src, tgt) in enumerate(val_iter):
        src = src.to(device)  # 将源序列移动到设备
        tgt = tgt.to(device)  # 将目标序列移动到设备

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

        # 创建掩码
        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:, :]  # 目标输出（去掉第一个 token）
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))  # 计算损失
        losses += loss.item()  # 累加损失

    return losses / len(val_iter)  # 返回平均损失

开始训练

在准备好必要的类和函数后，我们可以开始训练模型了。训练所需的时间可能会因为计算能力、参数设置和数据集大小等因素而有很大差异。

当我使用JParaCrawl完整的句子列表（每种语言大约有590万句）进行训练时，使用单个NVIDIA GeForce RTX 3060 GPU，每个epoch大约需要2小时。

以下是代码：

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS + 1)):
    start_time = time.time()
    train_loss = train_epoch(transformer, train_iter, optimizer)
    end_time = time.time()
    print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Epoch time = {(end_time - start_time):.3f}s")

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

首先，我们创建一些函数来翻译新句子，包括获取日语句子、分词、转换为张量、推理，然后将结果解码回句子，但这次是英文。

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    使用贪婪解码法进行序列生成。

    参数：
    - model: 训练好的 Transformer 模型。
    - src: 源序列张量，形状为 (src_seq_len, 1)。
    - src_mask: 源序列掩码张量，形状为 (src_seq_len, src_seq_len)。
    - max_len: 生成序列的最大长度。
    - start_symbol: 目标序列的起始符号索引。

    返回：
    - 生成的目标序列张量，形状为 (tgt_seq_len, 1)。
    """
    src = src.to(device)  # 将源序列移动到设备
    src_mask = src_mask.to(device)  # 将源序列掩码移动到设备
    memory = model.encode(src, src_mask)  # 编码源序列得到 memory
    ys = torch.ones(1, 1).fill_(start_symbol).type(torch.long).to(device)  # 初始化目标序列张量，起始符号为 start_symbol

    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)  # memory 掩码
        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: 训练好的 Transformer 模型。
    - src: 源语言句子（字符串）。
    - src_vocab: 源语言词汇表对象。
    - tgt_vocab: 目标语言词汇表对象。
    - src_tokenizer: 源语言分词器。

    返回：
    - 目标语言翻译后的句子（字符串）。
    """
    model.eval()  # 设置模型为评估模式

    # 对源语言句子进行分词并转换为索引
    tokens = [BOS_IDX] + [src_vocab.stoi[tok] for tok in src_tokenizer.encode(src, out_type=str)] + [EOS_IDX]
    num_tokens = len(tokens)
    src = torch.LongTensor(tokens).reshape(num_tokens, 1)  # 将源序列转换为张量
    src_mask = torch.zeros(num_tokens, num_tokens).type(torch.bool)  # 创建源序列掩码

    # 使用贪婪解码法生成目标序列
    tgt_tokens = greedy_decode(model, src, src_mask, max_len=num_tokens + 5, start_symbol=BOS_IDX).flatten()
    
    # 将目标序列索引转换为词，并去掉起始和结束符号
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

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

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

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

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

import pickle

# 保存 en_vocab
with open('en_vocab.pkl', 'wb') as file:
    pickle.dump(en_vocab, file)

# 保存 ja_vocab
with open('ja_vocab.pkl', 'wb') as file:
    pickle.dump(ja_vocab, file)

接下来，我们可以使用PyTorch的保存和加载功能保存模型以供以后使用。一般来说，有两种方式保存模型，取决于我们以后想要如何使用它们。第一种方式是仅用于推理，我们可以稍后加载模型并用它从日语翻译到英语。

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

第二种方式也是用于推理，但同时也用于我们想要稍后加载模型，并希望继续训练时。

# 保存模型和检查点以便稍后继续训练
torch.save({
    'epoch': NUM_EPOCHS,
    'model_state_dict': transformer.state_dict(),
    'optimizer_state_dict': optimizer.state_dict(),
    'loss': train_loss,
}, 'model_checkpoint.tar')

通过这些步骤，我们详细介绍了如何准备数据、构建模型、训练模型、评估模型以及保存模型。希望通过这篇博客，读者能够理解并实现一个日文-中文机器翻译模型。