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

基于Transformer&PyTorch的机器翻译模型

在本教程中，我们将使用Jupyter Notebook、PyTorch、Torchtext和SentencePiece构建一个日语到中文的机器翻译模型。

[2]:

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，请在你自己的电脑上尝试运行这一套代码

[3]:

device

[3]:

device(type='cpu')

获取并行数据集

在本教程中，我们将使用从JParaCrawl下载的日语-中文平行数据集。JParaCrawl是由NTT创建的“最大的公开可用的日语-英语平行语料库”，主要通过网络爬虫自动对齐平行句子创建而成。

[4]:

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中共有5973071个句子。然而，为了学习目的，通常建议先对数据进行采样，确保一切工作正常，然后再使用所有数据，以节省时间。

这里是数据集中的一对示例句子：

[5]:

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

We can also use different parallel datasets to follow along with this article, just make sure that we can process the data into the two lists of strings as shown above, containing the Japanese and English sentences.

准备分词器

与英语或其他字母语言不同，日语句子中没有空格来分隔单词。我们可以使用JParaCrawl提供的基于SentencePiece创建的分词器，分别用于日语和英语

[6]:

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

这些分词器将帮助我们将文本数据处理成模型可以接受的格式，用于训练我们的机器翻译模型。

加载了分词器之后，你可以通过执行以下代码进行测试。

[7]:

en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type='str')

[7]:

['▁All',
 '▁residents',
 '▁aged',
 '▁20',
 '▁to',
 '▁59',
 '▁years',
 '▁who',
 '▁live',
 '▁in',
 '▁Japan',
 '▁must',
 '▁enroll',
 '▁in',
 '▁public',
 '▁pension',
 '▁system',
 '.']

[8]:

ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。", out_type='str')

[8]:

['▁',
 '年',
 '金',
 '▁日本',
 'に住んでいる',
 '20',
 '歳',
 '~',
 '60',
 '歳の',
 '全ての',
 '人は',
 '、',
 '公的',
 '年',
 '金',
 '制度',
 'に',
 '加入',
 'しなければなりません',
 '。']

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

利用加载的分词器和原始句子，我们接着从TorchText中导入词汇表对象。这个过程的耗时取决于数据集的大小和计算能力，可能需要几秒或几分钟。不同的分词器也会影响构建词汇表所需的时间，我尝试了几种日语分词器，但发现SentencePiece对我来说既有效又快速。

[9]:

def build_vocab(sentences, tokenizer):
    # 创建一个Counter对象，用于统计每个token的出现次数
    counter = Counter()

    # 遍历所有的句子
    for sentence in sentences:
        # 使用分词器对每个句子进行分词，并转换为字符串类型l
        # 然后更新counter对象，统计每个token的出现次数
        counter.update(tokenizer.encode(sentence, out_type=str))

    # 使用Counter对象创建Vocab对象，同时定义特殊token
    # '<unk>'代表未知token，'<pad>'用于序列填充，'<bos>'和'<eos>'分别代表序列的开始和结束
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])
​
# 假设trainja和ja_tokenizer分别是训练集的日语句子和对应的分词器
ja_vocab = build_vocab(trainja, ja_tokenizer)
​
# 假设trainen和en_tokenizer分别是训练集的英语句子和对应的分词器
en_vocab = build_vocab(trainen, en_tokenizer)
​

当我们有了词汇表对象之后，可以利用词汇表和分词器对象为我们的训练数据构建张量。

[10]:

def data_process(ja, en):
    # 初始化一个空列表来存储处理后的数据对
    data = []
​
    # 使用zip函数遍历ja和en的对应元素
    for (raw_ja, raw_en) in zip(ja, en):
        # 对每个原始的ja和en句子进行处理
        # 使用ja_tokenizer将ja句子分词，并将结果编码为整数序列
        ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\\n"), out_type=str)], 
                                  dtype=torch.long)  # 注意：这里假设ja_vocab是之前构建的词汇表
​
        # 同理，使用en_tokenizer将en句子分词并编码
        en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("\\n"), out_type=str)], 
                                  dtype=torch.long)  # 注意：这里同样假设en_vocab是之前构建的词汇表
​
        # 将处理后的ja和en张量作为元组添加到data列表中
        data.append((ja_tensor_, en_tensor_))
​
    # 返回处理后的数据列表
    return data
​
# 假设trainja和trainen分别是训练集的日语和英语句子列表
train_data = data_process(trainja, trainen)

​

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

在这里，我将批量大小（BATCH_SIZE）设置为8，以避免“cuda out of memory”错误。实际的批量大小可以根据你的机器内存容量、数据集大小等因素进行调整。根据PyTorch教程使用Multi30k德英数据集的做法，批量大小设置为128

[11]:

# 设置批处理大小为8
BATCH_SIZE = 8
# 设置填充索引为日语词汇表中的<pad>对应的索引
PAD_IDX = ja_vocab['<pad>']
# 设置开始符号索引为日语词汇表中的<bos>对应的索引
BOS_IDX = ja_vocab['<bos>']
# 设置结束符号索引为日语词汇表中的<eos>对应的索引
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
# 使用DataLoader加载训练数据，设置批处理大小为8，打乱数据顺序，并使用generate_batch函数作为collate_fn参数
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是“注意力就是你所需要的”论文中介绍的一个Seq2Seq模型，用于解决机器翻译任务。转换器模型由编码器和解码器块组成，每个块包含固定数量的层。

[12]:

f

rom 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__()
        # 定义Transformer编码器层
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=NHEAD,
                                                dim_feedforward=dim_feedforward)
        # 创建Transformer编码器，包含多个编码器层
        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)
        # 创建Transformer解码器，包含多个解码器层
        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)
​

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

[ ]:

import torch
import torch.nn as nn
import math
​
# PositionalEncoding 类用于为输入序列添加位置信息
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: torch.Tensor):
        # 将词嵌入与位置编码相加，并应用丢弃层
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0),:])
​
# TokenEmbedding 类用于将输入的词索引转换为词向量
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: torch.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

Define model parameters and instantiate model. 这里我们服务器实在是计算能力有限，按照以下配置可以训练但是效果应该是不行的。如果想要看到训练的效果请使用你自己的带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)
​
# 将模型参数移动到指定设备上
transformer = transformer.to(device)
​
# 实例化交叉熵损失函数，并指定忽略索引
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)
​
# 实例化Adam优化器，并指定学习率、动量和eps参数
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
  with torch.no_grad():
    for idx, (src, tgt) in (enumerate(valid_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)
​
      # 前馈计算
​

开始训练

最后，在准备好必要的类和函数之后，我们就可以训练我们的模型了。这是不言而喻的，但完成训练所需的时间可能会因计算能力、参数和数据集大小等因素的不同而有很大差异。

当我使用JParaCrawl的完整句子列表训练模型时，每种语言大约有590万个句子，使用一个NVIDIA GeForce RTX 3070 GPU每个大约需要5个小时。

[ ]:

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}, "
          f"Epoch time = {(end_time - start_time):.3f}s"))

  0%|          | 0/16 [00:00<?, ?it/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(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)

​

[ ]:

trainen.pop(5)

[ ]:

trainja.pop(5)

[ ]:

​

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

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

[ ]:

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

Conclusion

That’s it!