Japanese-Chinese Machine Translation Model with Transformer

基于Transformer实现机器翻译（日译中）

Transformer模型简介

Transformer模型是一种用于处理序列数据的深度学习模型，特别是在自然语言处理领域中取得了显著的成就。它由Vaswani等人在2017年提出，主要解决了传统循环神经网络（RNNs）和长短期记忆网络（LSTMs）在处理长距离依赖问题上的局限性.

基于Transformer实现机器翻译（日译中）

导入环境

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')

我们将使用从JParaCrawl下载的日英平行数据集！[http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl]。JParaCrawl被描述为由NTT创建的“最大的公开可用的英日平行语料库”。该语料库主要通过网络爬取并自动对齐平行句子而创建。

df = pd.read_csv('./zh-ja/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条。然而，出于学习目的，通常建议先对数据进行抽样，确保一切按预期运行，然后再使用全部数据，以节省时间。

准备分词器

与英语或其他字母语言不同，日语句子中不包含空格来分隔单词。我们可以使用由JParaCrawl提供的分词器，这些分词器使用了SentencePiece来处理日语和英语。

en_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='enja_spm_models/spm.ja.nopretok.model')

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

使用分词器和原始句子，从 TorchText 导入词汇表对象。这个过程可能需要几秒钟或几分钟，取决于数据集大小和计算能力。不同的分词器也会影响构建词汇表所需的时间。

def build_vocab(sentences, tokenizer):
    """
    使用分词器从句子列表构建词汇表。

    Args:
    - sentences (list): 要构建词汇表的句子列表（字符串）。
    - tokenizer (Tokenizer): 用于分词句子的分词器对象。

    Returns:
    - 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):
    """
    将日语和英语句子处理成 Torch 张量形式。

    Args:
    - ja (list): 日语句子列表。
    - en (list): 英语句子列表。

    Returns:
    - list: 包含元组 (日语张量, 英语张量) 的数据列表。
    """
    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 out of memory” 错误，但实际上这取决于诸如机器内存容量、数据大小等各种因素，因此根据需要随意更改批量大小（注意：PyTorch教程在使用Multi30k德英数据集时将批量大小设置为128）。

BATCH_SIZE = 8  # 批量大小设置为 8
PAD_IDX = ja_vocab['<pad>']  # 日语词汇表中的 <pad> 符号的索引
BOS_IDX = ja_vocab['<bos>']  # 日语词汇表中的 <bos> 符号的索引
EOS_IDX = ja_vocab['<eos>']  # 日语词汇表中的 <eos> 符号的索引

def generate_batch(data_batch):
    """
    生成一个批次的数据，包括日语和英语句子的处理和填充。

    Args:
    - data_batch (list): 包含元组 (日语张量, 英语张量) 的数据批次列表。

    Returns:
    - tuple: 包含填充后的日语批张量和英语批张量的元组。
    """
    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)

序列到序列Transformer模型

下面的一些代码和文本解释（以斜体表示）摘自原始的PyTorch教程[https://pytorch.org/tutorials/beginner/translation_transformer.html]。

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)

文本标记通过使用标记嵌入来表示。在标记嵌入中添加位置编码，以引入单词顺序的概念。

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，表示单词在句子中的位置
        pos = torch.arange(0, maxlen).reshape(maxlen, 1)
        # 创建一个形状为 (maxlen, emb_size) 的零张量 pos_embedding
        pos_embedding = torch.zeros((maxlen, emb_size))
        # 根据位置和分母计算正弦和余弦值，并存储到 pos_embedding 中
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        # 在倒数第二个维度上添加一个维度，形状变为 (maxlen, emb_size, 1)
        pos_embedding = pos_embedding.unsqueeze(-2)

        # 定义一个 Dropout 层
        self.dropout = nn.Dropout(dropout)
        # 将 pos_embedding 注册为不可训练的缓冲区
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        # 在输入的 token_embedding 上应用 Dropout，并加上位置编码 pos_embedding
        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):
        # 输入 tokens 是一个整数张量，代表词汇索引，通过嵌入层得到词嵌入向量
        # 然后乘以 sqrt(emb_size) 进行缩放，返回嵌入后的向量表示
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建一个后续词掩码，防止目标词注意到其后续的词。我们还创建掩码，用于屏蔽源语言和目标语言的填充标记。

import torch

def generate_square_subsequent_mask(sz):
    # 生成一个上三角矩阵，主对角线及以下的元素为True，其余为False
    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))
    
    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)
    
    # 创建源序列的掩码（通常不需要，因为Transformer的自注意力机制已经处理）
    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的时候，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)
 
# 将模型移动到指定的设备（例如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, :]
 
      # 生成源序列和目标序列的mask
      src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
 
      # 模型前向传播计算logits
      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, :]
 
    # 生成源序列和目标序列的mask
    src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
 
    # 模型前向传播计算logits
    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)

开始训练

最终，在准备好必要的类和函数后，我们准备开始训练我们的模型。毫无疑问，训练完成所需的时间会因计算能力、模型参数和数据集大小等因素而有很大差异。
当我使用来自JParaCrawl的完整句子列表进行训练时，每种语言大约有590万句子，在单个NVIDIA GeForce RTX 3070 GPU上，每个epoch大约需要5小时。

import tqdm
import time

# 迭代训练多个epoch
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS + 1)):
    start_time = time.time()  # 记录当前epoch开始时间
    train_loss = train_epoch(transformer, train_iter, optimizer)  # 执行一个epoch的训练
    end_time = time.time()  # 记录当前epoch结束时间
    
    # 打印当前epoch的训练损失和epoch所花费的时间
    print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Epoch time = {(end_time - start_time):.3f}s")

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

首先，我们需要创建一系列函数来实现句子的翻译，包括以下步骤：获取日语句子、分词、转换为张量、推断（inference），然后将结果解码回一个英文句子。

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    # 将输入张量和掩码转移到指定设备上（通常是GPU）
    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 中
        ys = torch.cat([ys,
                        torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        
        # 如果生成了结束符号（EOS_IDX），则终止生成过程
        if next_word == EOS_IDX:
            break
    
    return ys


def translate(model, src, src_vocab, tgt_vocab, src_tokenizer):
    model.eval()  # 设置模型为评估模式，这通常会关闭 Dropout 等层

    # 对源语言进行标记化，并在两端添加起始符号（BOS_IDX）和结束符号（EOS_IDX）
    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()
    
    # 将标记化的目标语言序列转换为可读的字符串，并移除特殊标记（如<BOS>和<EOS>）
    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)
trainen.pop(5)

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

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

import pickle

# 打开一个文件，用于存储数据
file = open('en_vocab.pkl', 'wb')
# 将信息存储到文件中
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')

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

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