使用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')
获取并行数据集
在本教程中，我们将使用从JParaCrawl下载的日语-中文平行数据集。JParaCrawl是由NTT创建的“最大的公开可用的日语-英语平行语料库”，主要通过网络爬虫自动对齐平行句子创建而成。
df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainen = df[2].values.tolist()
trainja = df[3].values.tolist()
在导入所有的日语和其对应的中文数据之后，我删除了数据集中的最后一条数据，因为它有一个缺失值。总共，在trainen和trainja中共有5973071个句子。然而，为了学习目的，通常建议先对数据进行采样，确保一切工作正常，然后再使用所有数据，以节省时间。
准备数据集示例
这里是数据集中的一对示例句子：
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 ...
准备分词器
与英语或其他字母语言不同，日语句子中没有空格来分隔单词。我们可以使用JParaCrawl提供的基于SentencePiece创建的分词器，分别用于日语和英语
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')
构建TorchText的词汇表对象并将句子转换为Torch张量
利用加载的分词器和原始句子，我们接着从TorchText中导入词汇表对象。这个过程的耗时取决于数据集的大小和计算能力，可能需要几秒或几分钟。不同的分词器也会影响构建词汇表所需的时间，我尝试了几种日语分词器，但发现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): # 遍历每一对日语和英语句子
# 将日语句子编码为张量（tensor）
ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
dtype=torch.long)
# 将英语句子编码为张量（tensor）
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) # 处理训练数据集
通过以上步骤，我们已经准备好了日语和英语的训练数据集，可以用于后续的Transformer模型训练。接下来，我们将继续实现模型部分。
创建用于训练迭代的DataLoader对象
在这里，我将批量大小（BATCH_SIZE）设置为8，以避免“cuda out of memory”错误。实际的批量大小可以根据你的机器内存容量、数据集大小等因素进行调整。根据PyTorch教程使用Multi30k德英数据集的做法，批量大小设置为128
BATCH_SIZE = 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: # 遍历每个数据批次中的日语和英语张量对
# 在日语句子的开头和结尾添加 <bos> 和 <eos> 标记，然后将其拼接成张量
ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
# 在英语句子的开头和结尾添加 <bos> 和 <eos> 标记，然后将其拼接成张量
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 加载训练数据集，并设置批量大小、数据打乱、以及使用自定义的 generate_batch 函数进行数据处理
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
shuffle=True, collate_fn=generate_batch)
通过以上步骤，我们创建了一个用于训练的DataLoader对象 train_iter，它可以按批次提供填充后的日语和英语数据。接下来，可以将该对象用于Transformer模型的训练过程中。
这里是来自PyTorch官方教程的代码和相关文本解释，介绍了Transformer模型作为Seq2Seq模型用于机器翻译任务的实现过程。以下是对这部分内容的翻译和解释：
Transformer模型介绍
Transformer是一种Seq2Seq模型，最初在“Attention is all you need”论文中提出，用于解决机器翻译任务。Transformer模型由编码器和解码器组成，每个都包含固定数量的层。
编码器通过一系列的多头注意力和前馈网络层处理输入序列。编码器的输出称为“内存”，将其与目标张量一起传递给解码器。编码器和解码器通过教师强制（teacher forcing）技术进行端到端训练。
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,
t
文本标记通过使用标记嵌入来表示。位置编码通过添加到标记嵌入中引入了单词顺序的概念。
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), :])
以上代码和解释介绍了Transformer模型的主要组件，包括编码器、解码器、位置编码和标记嵌入的实现。这些组件将被用于构建Seq2Seq Transformer模型，用于日语到中文的机器翻译任务。
生成子序列遮罩和填充遮罩
def generate_square_subsequent_mask(sz):
# 生成一个对角线以下为1，对角线及以上为0的矩阵（上三角矩阵）
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)
# 生成源语言和目标语言的填充遮罩（PAD遮罩）
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
generate_square_subsequent_mask(sz): 生成一个对角线以下为1，对角线及以上为0的矩阵，用于在解码器中防止目标词汇关注到后续的词汇。create_mask(src, tgt): 根据源语言和目标语言序列，生成用于Transformer模型的各种遮罩。包括目标语言序列的子序列遮罩（用于解码器的自注意力机制），源语言序列的全零遮罩（用于编码器的自注意力机制），以及源语言和目标语言的填充遮罩（用于忽略填充位置的注意力计算）。
SRC_VOCAB_SIZE = len(ja_vocab) # 源语言词汇表大小 TGT_VOCAB_SIZE = len(en_vocab) # 目标语言词汇表大小 EMB_SIZE = 512 # 词嵌入的维度大小 NHEAD = 8 # 多头注意力机制的头数 FFN_HID_DIM = 512 # FeedForward层的隐藏层维度 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) # 初始化模型参数（Xavier初始化） for p in transformer.parameters(): if p.dim() > 1: nn.init.xavier_uniform_(p) # 将模型移动到设备（如GPU） transformer = transformer.to(device) # 损失函数：交叉熵损失，忽略填充标记PAD_IDX 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 )
定义了模型的参数，包括词汇表大小、词嵌入维度、多头注意力的头数、隐藏层维度等。
创建了Seq2SeqTransformer模型实例，并对其参数进行了Xavier初始化。
将模型移动到GPU（如果可用）。
定义了交叉熵损失函数和Adam优化器。
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) # 返回平均损失
这些函数完成了整个训练和评估过程，其中使用了之前定义的模型、损失函数、优化器以及生成的各种遮罩。
最终，在准备好必要的类和函数之后，我们可以开始训练我们的模型了。毫无疑问，训练完成所需的时间可能会因计算能力、参数设置和数据集大小等因素而有很大差异。
当我使用 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"))尝试使用训练好的模型翻译一句日语句子
首先，我们创建函数来翻译新的句子，包括获取日语句子、分词、转换为张量、推理，然后将结果解码回一个英语句子。
3.5
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) # 将分词结果转换为张量，添加维度为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() # 将生成的目标语言序列张量转换为词汇，并移除起始符号和终止符号 translation = " ".join([tgt_vocab.itos[tok] for tok in tgt_tokens]).replace("<bos>", "").replace("<eos>", "") return translation然后，我们只需调用翻译函数并传递所需的参数。
translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, en_vocab, ja_tokenizer)保存词汇表和训练好的模型：
最后，在训练完成后，我们将首先保存词汇表对象（en_vocab 和 ja_vocab）使用 Pickle。
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 的保存和加载函数保存模型以供以后使用。一般来说，根据以后的使用目的，有两种保存模型的方法。第一种适用于推理阶段，我们可以加载模型并用它将日语翻译成英语。
# 保存仅用于推理的模型 torch.save(transformer.state_dict(), 'inference_model')第二种方法也适用于推理，同时也适用于之后想要恢复训练的情况。
# 保存模型和检查点以便稍后恢复训练 torch.save({ 'epoch': NUM_EPOCHS, 'model_state_dict': transformer.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': train_loss, }, 'model_checkpoint.tar')
评论区