基于Transformer的机器翻译

构建一个基于Transformer和PyTorch的日文-中文机器翻译模型

在本教程中，我们将使用Jupyter Notebook、PyTorch、Torchtext和SentencePiece构建一个基于Transformer的日文-中文机器翻译模型。
在AutoDL平台租借GPU更方便

导入所需的包

首先，确保在系统中安装了以下包。如果发现某些包缺失，要确保安装它们。
在AutoDL的JupyterLab终端用pip命令下载

以下是要导入的库：

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

查看是否使用了GPU：

获取平行数据集（Get the parallel dataset）

在这次实验中，我们将使用从JParaCrawl ![http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl] 下载由NTT创建的日英平行数据集，该数据集是“最大规模的公开可用的英日平行语料库，主要通过网络爬取并自动对齐平行句子。”在链接中可以查看相关论文。
将数据集分成训练集（前70000）和验证集（剩下的）。并查看英文和日文是否对应。

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

结果对应：

在导入所有的日语和它们的英语对应数据后，删除了数据集中的最后一条数据，因为它存在缺失值。总共，训练集中日语和英语句子的数量为5,973,071条。然而，为了学习目的，通常建议对数据进行抽样，并确保一切按预期工作，然后再一次性使用所有数据，以节省时间。
下面是训练集中包含的一个句子示例。

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

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

准备分词器

与英语或其他字母语言不同，日语句子不包含空格来分隔单词（和中文类似，亚洲语种特点之一）。我们可以使用JParaCrawl提供的分词器，它们分别用于日语和英语，使用了SentencePiece。可以通过访问JParaCrawl网站下载它们。

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

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

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)

ja_tokenizer.encode("不同的应用场景对存储容量的要求也各不相同。", out_type=str)

en_tokenizer.encode("不同的应用场景对存储容量的要求也各不相同。", out_type=str)

日文分词器分词中文效果比英文分词器好一点。

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

使用分词器和原始句子，接着从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)
valid_data = data_process(validja, validen)

创建用于训练过程中迭代的DataLoader对象

这里，将BATCH_SIZE设置为8，以避免“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))
  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)
valid_iter = DataLoader(valid_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

序列到序列的Transformer模型

接下来的几段代码和文字说明取自 [https://pytorch.org/tutorials/beginner/translation_transformer.html]。除了将BATCH_SIZE和单词de_vocab改为ja_vocab之外，没有做任何更改。
Transformer是一种在论文“Attention is all you need”中引入的Seq2Seq模型，用于解决机器翻译任务。Transformer模型由编码器和解码器块组成，每个块包含固定数量的层。
编码器通过一系列多头注意力和前馈网络层处理输入序列。编码器的输出，称为记忆，被传递给解码器和目标张量。编码器和解码器通过使用 teacher forcing technique（教师强制技术：每次不使用上一个state的输出作为下一个state的输入，而是直接使用训练数据的标准答案(ground truth)的对应上一项作为下一个state的输入。）进行端到端训练。

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 = 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):
        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)
    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)
  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设置为128。

SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(en_vocab)
EMB_SIZE = 128
NHEAD = 8
FFN_HID_DIM = 128
BATCH_SIZE = 8
NUM_ENCODER_LAYERS = 3
NUM_DECODER_LAYERS = 3
NUM_EPOCHS = 8
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)

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)

开始训练模型

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

for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
  start_time = time.time()
  train_loss = train_epoch(transformer, train_iter, optimizer)
  valid_loss = train_epoch(transformer, valid_iter, optimizer)
  end_time = time.time()
  print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Valid loss: {valid_loss:.3f}"
          f"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()
    return " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")

然后，可以直接调用 translate_sentence 函数并传递必要的参数。

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

trainen.pop(5)

trainja.pop(5)

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

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

# 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文件和inference_model文件已被存入

通过上述步骤，我们已经完成了一个基于Transformer的日中机器翻译模型。通过不断训练和调整参数，模型性能会有所提升。

总结

基于Transformer的机器翻译模型已经成为目前最先进的方法之一，为自然语言处理领域的发展带来了重要的推动力。
Transformer模型通过自注意力机制实现了全局上下文的建模，使得模型能够并行处理输入序列，大大提高了训练和推理效率。相比传统的循环神经网络（RNN）和卷积神经网络（CNN），Transformer能够更有效地捕捉长距离的依赖关系，从而提升了翻译质量。
总之，基于Transformer的机器翻译模型不仅在学术界取得了显著成就，也在工业界得到了广泛应用，为机器翻译和自然语言处理领域带来了重要的技术进步和应用前景。