基于transformer的日译中机器翻译模型

导入所需要的包

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

获取并行数据集

我们将使用从JParaCrawl![http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl]]下载的日语-英语并行数据集，该数据集被描述为“NTT创建的最大的公开可用的英语-日语并行语料库”。它主要是通过抓取网络并自动对齐平行句子创建的。

# 导入pandas库
import pandas as pd

# 使用pandas读取文件，文件路径为'./zh-ja/zh-ja.bicleaner05.txt'
# 文件使用制表符（\t）作为分隔符，使用python引擎来解析文件，无表头
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 将第三列（索引为2）的值转换为列表，存储在变量trainen中
trainen = df[2].values.tolist()#[:10000]

# 将第四列（索引为3）的值转换为列表，存储在变量trainja中
trainja = df[3].values.tolist()#[:10000]

# 移除列表trainen和trainja中索引为5972的元素
# trainen.pop(5972)
# trainja.pop(5972)

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

准备标记器

与英语或其他按字母顺序排列的语言不同，日语句子不包含空格来分隔单词。我们可以使用JParaCrawl提供的标记器，它是使用sentencepece为日语和英语创建的，您可以访问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')

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

使用标记器和原始句子，然后构建从TorchText导入的Vocab对象。根据数据集的大小和计算能力，这个过程可能需要几秒钟或几分钟。不同的标记器也会影响构建词汇所需的时间，我尝试了其他几个日语标记器，但sensenepece似乎工作得很好，对我来说足够快。

# 导入Counter类用于统计词频
from collections import Counter
# 导入Vocab类用于创建词汇表
from some_module import Vocab  # 假设Vocab类是从some_module导入的

# 定义一个函数，用于构建词汇表
def build_vocab(sentences, tokenizer):
    # 创建一个Counter对象，用于统计词频
    counter = Counter()
    # 遍历所有句子
    for sentence in sentences:
        # 使用分词器对句子进行编码，并更新词频统计
        counter.update(tokenizer.encode(sentence, out_type=str))
    # 返回一个Vocab对象，包含词频统计和特殊标记
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])

# 使用训练集中的日语句子和日语分词器构建日语词汇表
ja_vocab = build_vocab(trainja, ja_tokenizer)

# 使用训练集中的英语句子和英语分词器构建英语词汇表
en_vocab = build_vocab(trainen, en_tokenizer)

# 导入torch库，用于处理张量
import torch

# 定义一个函数，用于处理数据
def data_process(ja, en):
    # 创建一个空列表，用于存储处理后的数据
    data = []
    # 使用zip函数将日语和英语句子配对遍历
    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列表中
        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  # 每个批次的大小为8
PAD_IDX = ja_vocab['<pad>']  # 获取'<pad>'的索引
BOS_IDX = ja_vocab['<bos>']  # 获取'<bos>'的索引
EOS_IDX = ja_vocab['<eos>']  # 获取'<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加载训练数据
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)  # 设置批次大小，打乱数据，并使用generate_batch函数处理数据

Sequence-to-sequence变压器

接下来的代码和文本解释(以斜体书写)来自原始的PyTorch教程[https://pytorch.org/tutorials/beginner/translation_transformer.html]]。除了BATCH_SIZE和单词de_vocab被更改为ja_vocab之外，我没有做任何更改。

Transformer是在“Attention is all you need”论文中提出的用于解决机器翻译任务的Seq2Seq模型。变压器模型由编码器和解码器块组成，每个块包含固定数量的层。

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

from torch.nn import (TransformerEncoder, TransformerDecoder,
                      TransformerEncoderLayer, TransformerDecoderLayer)

# 定义序列到序列的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__()
        
        # 定义编码器层
        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)

import torch
import torch.nn as nn
import math
from torch import Tensor

# 定义位置编码类
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)
        
        # 定义dropout层
        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):
        # 返回输入tokens的词嵌入表示，并乘以sqrt(emb_size)以缩放
        return self.embedding(tokens.long()) * math.sqrt(self.emb_size)

import torch

# 生成Transformer中使用的前瞻遮挡（用于解码器中）
def generate_square_subsequent_mask(sz):
    # 创建一个大小为(sz, sz)的上三角矩阵
    mask = (torch.triu(torch.ones((sz, sz), device=device)) == 1).transpose(0, 1)
    # 将矩阵转换为float类型，并将为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)
    # 创建源序列的填充遮挡（全零矩阵，掩盖掉填充位置）
    src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)

    # 创建源序列和目标序列的填充遮挡（用于DataLoader的collate_fn函数中）
    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  # 训练的轮数

# 创建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)

# 将模型移动到指定的设备（如GPU）
transformer = transformer.to(device)

# 定义损失函数（交叉熵损失函数）
loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)

# 定义优化器（Adam优化器）
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
        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
        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万个句子)的完整句子列表来训练模型时，使用单个NVIDIA GeForce RTX 3070 GPU，每个epoch大约需要5个小时。

代码如下:

# 使用tqdm显示训练进度，并输出每个epoch的训练损失和训练时间
for epoch in tqdm.tqdm(range(1, NUM_EPOCHS + 1)):
    start_time = time.time()  # 记录当前epoch开始的时间
    train_loss = train_epoch(transformer, train_iter, optimizer)  # 训练一个epoch
    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"))

试着用训练好的模型翻译一个日语句子

首先，我们创建翻译新句子的函数，包括获取日语句子、标记化、转换为张量、推理，然后将结果解码回句子，但这次是英语。

# 使用贪婪解码进行序列到序列的翻译
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)
    
    # 将分词后的文本转换为Tensor，并添加一个维度表示批次大小为1
    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>", "")

保存Vocab对象和训练好的模型

最后，在训练完成后，我们将首先使用Pickle保存Vocab对象(en_vocab和ja_vocab)。

import pickle

# 将英语词汇表保存到文件'en_vocab.pkl'
with open('en_vocab.pkl', 'wb') as file:
    pickle.dump(en_vocab, file)

# 将日语词汇表保存到文件'ja_vocab.pkl'
with open('ja_vocab.pkl', 'wb') as file:
    pickle.dump(ja_vocab, file)

最后，我们还可以使用PyTorch保存和加载函数保存模型以供以后使用。通常，有两种保存模型的方法，这取决于我们以后想要使用它们的目的。第一个仅用于推理，我们可以稍后加载模型并使用它从日语翻译成英语。

# save model for inference
torch.save(transformer.state_dict(), 'inference_model')

第二个也是用于推理的，但当我们稍后想要加载模型并想要恢复训练时也是如此。

import torch

# 保存模型和检查点以便稍后恢复训练
torch.save({
    'epoch': NUM_EPOCHS,                     # 保存当前训练的轮数（epoch）
    'model_state_dict': transformer.state_dict(),   # 保存模型的状态字典（包含模型的参数）
    'optimizer_state_dict': optimizer.state_dict(), # 保存优化器的状态字典（包含优化器的状态）
    'loss': train_loss,                       # 保存当前训练的损失值
}, 'model_checkpoint.tar')                    # 将所有信息保存到'model_checkpoint.tar'文件中