基于Transformer实现机器翻译

实验前言

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

一、Import required packages

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

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

device(type='cpu')

二、Get the parallel dataset

        在本实验中，我们将使用从JParaCrawl下载的日英平行数据集！http://www.kecl.ntt.co.jp/icl/lirg/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])  # 打印trainen列表中索引为500的元素
print(trainja[500])  # 打印trainja列表中索引为500的元素

我们还可以使用不同的并行数据集来遵循本文，只需确保我们可以将数据处理成两个字符串列表，如上所示，包含日语和英语句子。

三、Prepare the tokenizers

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

# 这是英文分词器（Tokenizer）的实例化
en_tokenizer = spm.SentencePieceProcessor(model_file='spm.en.nopretok.model')

# 这是日文分词器（Tokenizer）的实例化
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')

四、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):
    # 将日语句子进行分词，并将分词后的词汇转换为对应的索引值
    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)

五、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 的元素

九、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. 数据准备：

   • 收集并准备日文和中文平行语料库作为训练数据。
   • 划分数据集为训练集、验证集和测试集。
   • 执行数据预处理，例如分词、建立词典等。

2. 模型架构设计：

   • 实现Transformer模型的编码器和解码器结构。
   • 确定词嵌入维度、隐藏层大小、注意力头数等超参数。
   • 添加位置编码以帮助模型对序列进行建模。

3. 模型训练：

   • 在PyTorch中构建Transformer模型。
   • 定义损失函数，通常采用交叉熵损失函数。
   • 使用优化器如Adam或AdamW来训练模型。
   • 实现训练循环，并在训练集和验证集上进行训练和评估。

4. 推理和评估：

   • 使用训练好的模型进行推理，翻译目标语言句子。
   • 评估模型性能，例如BLEU分数等指标。

5. 部署和优化：

   • 将训练好的模型部署到生产环境中。
   • 可以进行模型微调、超参数调整等优化工作。

总的来说，构建日中机器翻译模型需要很多步骤和考虑事项，以确保最终模型能够在实际应用中达到良好的性能。