基于Transformer实现的机器翻译

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

使用 Jupyter Notebook、PyTorch、Torchtext 和 SentencePiece 的教程

环境及软件包

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

获取并行数据集

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

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

# 提取文件中第三列的数据，转换为列表，作为训练数据的英文部分
trainen = df[2].values.tolist() #[:10000] # 可以限制只取前10000条数据

# 提取文件中第四列的数据，转换为列表，作为训练数据的日文部分
trainja = df[3].values.tolist()

# trainen.pop(5972)
# trainja.pop(5972)

在导入所有日语和英语对应数据后，我删除了数据集中的最后一个数据，因为它缺少值。总的来说，trainen和trainja中的句子数量都是5973071，然而，出于学习目的，通常建议在一次使用所有数据之前，对数据进行采样，并确保一切正常，以节省时间。

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

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

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

准备标记器

与英语或其他按字母顺序排列的语言不同，日语句子不包含空格来分隔单词。我们可以使用JParaCrawl提供的标记器，该标记器是使用日语和英语的句子片段创建的，您可以访问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对句子进行编码，编码结果为字符串类型
encoded_sentence = en_tokenizer.encode(
    "All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", 
    out_type='str'
)
print(encoded_sentence)

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

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

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

构建TorchText Vocab对象并将句子转换为Torch张量

使用标记器和原始句子，我们构建从TorchText导入的Vocab对象。这个过程可能需要几秒钟或几分钟，这取决于我们数据集的大小和计算能力。不同的标记器也会影响构建vocab所需的时间，我尝试了其他几种日语标记器，但PensionePiece似乎对我来说足够好、足够快。

# 构建词汇表函数
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>'])

# 构建英文和日文的词汇表
en_vocab = build_vocab(trainen, en_tokenizer)
ja_vocab = build_vocab(trainja, ja_tokenizer)

在我们有了词汇表对象之后，我们可以使用vocab和标记器对象来为我们的训练数据构建张量。

# 数据处理函数，将日文和英文句子转换为张量对
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)

创建要在训练期间迭代的DataLoader对象

在这里，我将BATCH_SIZE设置为16，以防止“cuda内存不足”，但这取决于各种因素，如机器内存容量、数据大小等，所以可以根据您的需要随意更改批量大小（注意：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:
        # 将BOS符号添加到日文句子开头，将EOS符号添加到日文句子结尾
        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进行批量处理
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模型，用于解决机器翻译任务。转换器模型由编码器和解码器块组成，每个块包含固定数量的层。

编码器通过一系列多头注意和前馈网络层传播输入序列来处理输入序列。编码器的输出（称为存储器）与目标张量一起被馈送到解码器。编码器和解码器使用教师强制技术以端到端的方式进行训练。

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

        # 定义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):
        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)

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

# 定义位置编码类，继承自nn.Module
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)  # 偶数索引位置使用sin函数
        pos_embedding[:, 1::2] = torch.cos(pos * den)  # 奇数索引位置使用cos函数

        # 在最后一个维度上添加一个新的维度
        pos_embedding = pos_embedding.unsqueeze(-2)

        # 定义Dropout层
        self.dropout = nn.Dropout(dropout)

        # 注册位置编码为缓冲区，不会更新
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        # 在输入的token嵌入上添加位置编码，并应用Dropout
        return self.dropout(token_embedding + self.pos_embedding[:token_embedding.size(0), :])

# 定义TokenEmbedding类，继承自nn.Module
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):
        # 对输入的tokens进行嵌入，并乘以嵌入维度的平方根
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

我们创建一个后续单词掩码来阻止目标单词关注其后续单词。我们还创建掩码，用于掩码源和目标填充令牌。

# 生成用于Transformer解码器的遮蔽矩阵，保证解码时只能看到之前已生成的部分
def generate_square_subsequent_mask(sz):
    # 创建一个上三角矩阵，并转置为下三角矩阵
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将矩阵中的1替换为负无穷，0替换为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)
    
    # 创建源语言的全零遮蔽矩阵，因为Transformer编码器不需要遮蔽
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)

    # 创建源语言和目标语言的填充遮蔽张量
    src_padding_mask = (src == PAD_IDX).transpose(0, 1)  # 将PAD符号位置设为True
    tgt_padding_mask = (tgt == PAD_IDX).transpose(0, 1)  # 将PAD符号位置设为True

    # 返回源语言遮蔽矩阵、目标语言遮蔽矩阵、源语言填充遮蔽张量、目标语言填充遮蔽张量
    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)

# 定义优化器
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)

        # 去除目标序列的最后一个token作为解码器输入
        tgt_input = tgt[:-1, :]

        # 创建源语言和目标语言的遮蔽张量
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

        # 前向传播，计算logits
        logits = model(src, tgt_input, src_mask, tgt_mask,
                       src_padding_mask, tgt_padding_mask, src_padding_mask)

        optimizer.zero_grad()

        # 去除目标序列的第一个token作为计算损失的真实目标
        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)

        # 去除目标序列的最后一个token作为解码器输入
        tgt_input = tgt[:-1, :]

        # 创建源语言和目标语言的遮蔽张量
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)

        # 前向传播，计算logits
        logits = model(src, tgt_input, src_mask, tgt_mask,
                       src_padding_mask, tgt_padding_mask, src_padding_mask)

        # 去除目标序列的第一个token作为计算损失的真实目标
        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万个句子，使用一个NVIDIA GeForce RTX 3070 GPU每个历元大约需要5个小时。

这是代码：

# 使用 tqdm.tqdm 进行训练过程的可视化
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的训练损失和训练时间
    print(f"Epoch: {epoch}, Train loss: {train_loss:.3f}, 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)
    
    # 初始化目标语言序列，起始符号为start_symbol
    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)
        
        # 解码生成下一个目标语言token
        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()
        
        # 将生成的下一个token添加到目标语言序列中
        ys = torch.cat([ys,
                        torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=0)
        
        # 如果生成了结束符号<eos>，则停止生成
        if next_word == EOS_IDX:
            break
    
    return ys

#翻译函数，将源语言句子翻译为目标语言句子
def translate(model, src, src_vocab, tgt_vocab, src_tokenizer):
    #将模型修改为评估模式
    model.eval()
    
    # 对源语言句子进行分词和编码，并加入起始符号<bos>和结束符号<eos>
    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()
    
    # 将生成的目标语言序列张量转换为字符串形式，去除<bos>和<eos>标记
    translated_sentence = " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "")
    
    return translated_sentence

然后，我们可以直接调用translate函数并传递所需的参数。

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

保存Vocab对象和训练的模型

最后，在训练完成后，我们将首先使用Pickle保存Vocab对象（en_Vocab和ja_Vocab）。

import pickle 

# 打开一个文件，用于存储数据
file = open('en_vocab.pkl', 'wb')  

# 使用pickle将en_vocab对象存储到文件中
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')

结束