基于Transformer和PyTorch构建中日机器翻译模型

提示：文章写完后，目录可以自动生成，如何生成可参考右边的帮助文档

---

实验前言

提示：这里可以添加本文要记录的大概内容：

    在当今信息时代，机器翻译技术的发展日新月异，特别是基于深度学习的方法已经成为推动翻译质量和效率提升的关键。本文将探讨如何利用Transformer模型和PyTorch框架，构建一个高效的中日双语翻译系统。Transformer模型以其优秀的并行性和注意力机制，极大地改善了长距离依赖问题，使得其在序列到序列任务中表现突出。我将从数据预处理、模型架构设计到训练和推断的全过程进行详细讨论，旨在为读者提供一种深入理解和实现现代机器翻译技术的实用指南。通过本文的学习，希望读者能够掌握构建和优化Transformer模型的关键步骤，为日后在实际应用中应对多样化的语言翻译挑战提供强有力的技术支持。实验步骤如下：

transformer原理

提示：以下是本篇文章正文内容，下面案例可供参考

一、Import required packages

    首先，让我们确保在系统中安装了以下软件包，如果发现某些软件包缺失，请确保安装它们。

import math  # 导入数学模块
import torchtext  # 导入torchtext库
import torch  # 导入PyTorch库
import torch.nn as nn  # 导入神经网络模块
from torch import Tensor  # 从PyTorch中导入张量
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  # 导入时间模块
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或CPU）
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 如果有GPU可用，打印出GPU名称
print(torch.cuda.get_device_name(0))  # 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

device

二、Get the parallel dataset

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

# 读取制表符分隔的文件，生成DataFrame
df = pd.read_csv('/root/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 将第三列（索引为2）的数据转换为列表，作为训练的英语句子
trainen = df[2].values.tolist()  # [:10000]

# 将第四列（索引为3）的数据转换为列表，作为训练的日语句子
trainja = df[3].values.tolist()  # [:10000]
# trainen.pop(5972)
# trainja.pop(5972)

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

Here is an example of sentence contained in the dataset.

# 打印训练数据中的第500个英语句子
print(trainen[500])

# 打印训练数据中的第500个日语句子
print(trainja[500])

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

三、Prepare the tokenizers

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

# 加载英语的SentencePiece模型，用于分词
en_tokenizer = spm.SentencePieceProcessor(model_file='/root/spm.en.nopretok.model')

# 加载日语的SentencePiece模型，用于分词
ja_tokenizer = spm.SentencePieceProcessor(model_file='/root/spm.ja.nopretok.model')

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

en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.")# 使用英语的SentencePiece

ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。")# 使用英语的SentencePiece模型

四、Build the TorchText Vocab objects and convert the sentences into Torch tensors

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

# 定义构建词汇表的函数
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 = []
    # 遍历每对日语和英语句子
    for (raw_ja, raw_en) in zip(ja, en):
        # 使用日语tokenizer将句子分词并转换为索引张量
        ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                                  dtype=torch.long)
        # 使用英语tokenizer将句子分词并转换为索引张量
        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)

五、Create the DataLoader object to be iterated during training

    在这里，我将 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)

六、Sequence-to-sequence Transformer

    接下来的几段代码和文本解释（用斜体表示）摘自原始的PyTorch教程https://pytorch.org/tutorials/beginner/translation_transformer.html。Transformer是一种Seq2Seq模型，引入于“Attention is all you need”论文中，用于解决机器翻译任务。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__()
        
        # 初始化Transformer编码器层
        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)
        
        # 线性层用于生成输出词汇表
        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))
        
        # 编码阶段：使用Transformer编码器
        memory = self.transformer_encoder(src_emb, src_mask, src_padding_mask)
        
        # 解码阶段：使用Transformer解码器
        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):
        # 编码阶段：仅使用Transformer编码器
        return self.transformer_encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        # 解码阶段：仅使用Transformer解码器
        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))
        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):
        # 将输入的词索引张量转换为词嵌入向量，并乘以一个缩放因子
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

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

def generate_square_subsequent_mask(sz):
    # 生成一个上三角矩阵，并转置为下三角矩阵
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(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)
    # 创建源语言的掩码（全为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

# 初始化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)

# 将模型移动到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):
    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(val_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)

七、Start training

    毫无疑问，训练完成所需的时间会因计算能力、参数设置和数据集大小等因素而有很大差异。

    当使用完整的JParaCrawl数据集（每种语言约590万句子）训练模型时，使用单个 NVIDIA GeForce RTX 3070 GPU 每个epoch大约需要5个小时。我这里用的是NVIDIA GeForce RTX 4090 D，总共需要1.5小时。

# 开始训练过程
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}, "
           f"Epoch time = {(end_time - start_time):.3f}s"))

八、Try translating a Japanese sentence using the trained model

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

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    src = src.to(device)  # 将输入的源语言序列移动到设备（GPU或CPU）
    src_mask = src_mask.to(device)  # 将掩码也移动到设备（GPU或CPU）
    
    # 编码源语言序列
    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()
    
    # 将目标语言序列转换为字符串并去除特殊符号
    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)
# 使用 Seq2Seq 模型进行翻译

（看得出来效果并不好）

trainen.pop(5)# 移除 trainen 列表中索引为 5 的元素

trainja.pop(5)# 移除 trainja 列表中索引为 5 的元素

九、Save the Vocab objects and trained model

    最后，在训练完成后，我们将首先使用Pickle保存Vocab对象（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')

总结

    使用Transformer和PyTorch构建日中机器翻译模型涉及几个关键步骤。以下是创建这种模型的一般步骤和考虑事项：

1，数据准备：

获取日语-英语平行语料库，如JParaCrawl。
清洗和分词日语和英语数据。

2，构建词汇表：

使用分词器（如SentencePiece）为日语和英语构建词汇表。

3，模型定义：

定义Transformer模型，包括编码器和解码器，以及相关的嵌入层、注意力机制和前馈网络。

4，模型训练：

设置训练参数，如批处理大小、学习率等。
使用训练数据训练Transformer模型。
监控训练过程中的损失和性能指标。

5，模型评估：

使用验证集评估模型的翻译质量，如计算BLEU分数。

6，模型推断：

定义推断函数，用训练好的模型进行日语到英语的翻译。

7，模型保存：

使用Pickle保存词汇表对象（en_vocab和ja_vocab）。
使用PyTorch的保存函数保存模型，以便后续推断或恢复训练使用。

8，后续应用：

可以加载保存的模型进行实时翻译服务或继续训练。

   这个流程涵盖了从数据准备到模型训练和保存的全过程，确保了能够建立一个日语-英语机器翻译模型并对其进行有效的管理和使用。希望能对读者在transformer方面的学习产生些帮助。