使用Transformer和PyTorch的日-中机器翻译模型使用--Jupyter、 Notebook、PyTorch、Torchtext和SentencePiece的教程

一、导入所需的包

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

# 导入数学库
import math
# 导入torchtext库
import torchtext         
# 导入torch库
import torch
# 导入torch.nn库并简写为nn
import torch.nn as nn
# 从torch库中导入Tensor类
from torch import Tensor
# 从torch.nn.utils.rnn库中导入pad_sequence函数
from torch.nn.utils.rnn import pad_sequence
# 从torch.utils.data库中导入DataLoader类
from torch.utils.data import DataLoader
# 从collections库中导入Counter类
from collections import Counter
# 从torchtext.vocab库中导入Vocab类
from torchtext.vocab import Vocab
# 从torch.nn库中导入TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer类
from torch.nn import TransformerEncoder, TransformerDecoder, TransformerEncoderLayer, TransformerDecoderLayer
# 导入io库
import io
# 导入time库
import time
# 导入pandas库并简写为pd
import pandas as pd
# 导入numpy库并简写为np
import numpy as np
# 导入pickle库
import pickle
# 导入tqdm库
import tqdm
# 导入sentencepiece库并简写为spm
import sentencepiece as spm
# 设置随机种子为0
torch.manual_seed(0)
# 判断是否有GPU，如果有则使用GPU，否则使用CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# 如果你有GPU，请在你自己的电脑上尝试运行这一套代码
# print(torch.cuda.get_device_name(0))

device

运行结果：

device(type='cpu')

二、获取平行数据集

在这个教程中，我们将使用从JParaCrawl下载的日英平行数据集！[http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl]，这被描述为“NTT创建的最大的可公开获取的英日平行语料库。它是通过大规模爬取网络并自动对齐平行句子而创建的。” 您也可以在这里看到论文。

 

# 导入pandas库
import pandas as pd

# 读取名为'zh-ja.bicleaner05.txt'的文件，使用制表符分隔，使用Python引擎，没有列名
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 将第3列的数据转换为列表，赋值给trainen
trainen = df[2].values.tolist()

# 将第4列的数据转换为列表，赋值给trainja
trainja = df[3].values.tolist()

# 注释掉的代码：从trainen中删除索引为5972的元素
# trainen.pop(5972)

# 注释掉的代码：从trainja中删除索引为5972的元素
# trainja.pop(5972)

在导入所有日语及其英文对应文本之后，我删除了数据集中的最后一条数据，因为它缺少数值。总共，在训练集（trainen）和训练日语集（trainja）中句子的数量为 5,973,071 条，然而，为了学习目的，通常建议对数据进行抽样，并确保一切按预期运行，然后再一次性使用所有数据，以节省时间。

以下是数据集中包含的一个句子的示例。

# 打印trainen列表中索引为500的元素
print(trainen[500])
# 打印trainja列表中索引为500的元素
print(trainja[500])

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

三、准备分词器

不像英语或其他字母语言，日语句子中不包含空格来分隔单词。我们可以使用由SentencePiece创建的JParaCrawl提供的分词器，适用于日语和英语，您可以访问JParaCrawl网站下载它们，或者点击这里。

# 创建英文分词器实例，加载英文模型文件
en_tokenizer = SentencePieceProcessor(model_file='enja_spm_models/spm.en.nopretok.model')

# 创建日文分词器实例，加载日文模型文件
ja_tokenizer = 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')

 

四、构建 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("
"), out_type=str)],
                            dtype=torch.long)
    # 对英文句子进行分词，并将分词结果转换为词汇表中对应的索引值，生成一个张量
    en_tensor_ = torch.tensor([en_vocab[token] for token in en_tokenizer.encode(raw_en.rstrip("
"), out_type=str)],
                            dtype=torch.long)
    # 将日文和英文张量组成一个元组，添加到数据列表中
    data.append((ja_tensor_, en_tensor_))
  # 返回处理后的数据列表
  return data

# 调用data_process函数，处理训练集数据，并将结果赋值给train_data变量
train_data = data_process(trainja, trainen)

 五、创建DataLoader对象以在训练期间进行迭代

这里，我将BATCH_SIZE设为16，以防止“cuda内存溢出”，但这取决于各种因素，比如您的机器内存容量、数据大小等，所以根据您的需要随时更改批量大小（注意：PyTorch的教程在Multi30k德英数据集上将批量大小设为128）。

# 设置每个批次的大小为8
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

# 创建一个数据迭代器，输入为训练数据，批次大小为8，打乱顺序，并使用generate_batch函数作为批处理函数
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

六、序列到序列变换器

接下来的一对代码和文本解释（用斜体书写）取自原始的PyTorch教程

[https://pytorch.org/tutorials/beginner/translation_transformer.html]。除了BATCH_SIZE和单词de_vocab被更改为ja_vocab之外，我没有进行任何更改。

Transformer是一种Seq2Seq模型，这种模型在“注意力机制就是一切”论文中被引入，用于解决机器翻译任务。Transformer模型包括一个编码器和一个解码器块，每个块都包含固定数量的层。

编码器通过传播输入序列，经过一系列Multi-head Attention和Feed forward网络层来处理输入序列。编码器的输出被称为内存，会与目标张量一起被传递给解码器。编码器和解码器会使用教师强制技术以端到端的方式进行训练。

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):
    # 初始化函数，输入参数为emb_size（嵌入维度）、dropout（丢弃率）和maxlen（最大长度，默认为5000）
    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)

        # 初始化dropout层
        self.dropout = nn.Dropout(dropout)
        # 注册位置嵌入矩阵为buffer
        self.register_buffer('pos_embedding', pos_embedding)

    # 前向传播函数，输入参数为token_embedding（词嵌入矩阵）
    def forward(self, token_embedding: Tensor):
        # 返回经过dropout处理的位置嵌入矩阵与词嵌入矩阵相加的结果
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0),:])

# 定义词嵌入类
class TokenEmbedding(nn.Module):
    # 初始化函数，输入参数为vocab_size（词汇表大小）和emb_size（嵌入维度）
    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

    # 前向传播函数，输入参数为tokens（词序列）
    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替换为负无穷，1替换为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)
    # 创建一个全零矩阵作为源序列的掩码
    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
# 初始化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或CPU）
transformer = transformer.to(device)

# 定义损失函数，忽略PAD_IDX位置的损失
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(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)

    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万句子）训练模型时，每个时代大约需要5个小时，使用单个 NVIDIA GeForce RTX 3070 GPU。

这里是代码：

# 使用tqdm库显示训练进度条
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS+1)):
    # 记录当前时间，用于计算每个epoch的耗时
    start_time = time.time()
    # 调用train_epoch函数进行一轮训练，并返回训练损失值
    train_loss = train_epoch(transformer, train_iter, optimizer)
    # 记录当前时间，用于计算每个epoch的耗时
    end_time = time.time()
    # 打印当前轮次、训练损失值和耗时信息
    print((f"Epoch: {epoch}, Train loss: {train_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.
    # 对源语言序列进行编码
    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列表中移除索引为5的元素
trainen.pop(5)

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

九、保存词汇表对象和训练模型

训练结束后，我们将首先使用Pickle保存Vocab对象（en_vocab和ja_vocab）。

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

The second one is for inference too, but also for when we want to load the model later, and want to resume the training

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