带有Transformer和PyTorch的日中机器翻译模型

Gavin299

已于 2024-06-27 12:48:26 修改

阅读量1.1k

点赞数 41

文章标签： transformer pytorch 机器翻译自然语言处理 python

于 2024-06-23 01:17:12 首次发布

本文链接：https://blog.csdn.net/Gavin299/article/details/139890777

版权

一个使用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')

device

结果：device(type='cpu')

二、获取并行数据集

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

# 读取名为'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(5927)

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

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

print(trainen[500])#打印列表中索引为500的元素
print(trainja[500])#打印列表中索引为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 ...

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

三、准备分词器

与英语或其他字母语言不同，日语句子不包含空格来分隔单词。我们可以使用由SentencePiece创建的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')

加载tokenizers后，您可以通过执行下面的代码来测试它们。

en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type='str')
# 使用英文分词器对句子进行编码，输出类型为字符串

实验结果：

['▁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词汇表对象并将句子转换为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)
# 使用训练集数据构建英文词汇表

在我们有了词汇表对象之后，我们可以使用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("
"), 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)

六、序列到序列Transformer

接下来的几个代码和文本解释（用斜体书写）取自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__()
        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):
    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: Tensor):
        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):
        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万句子的完整列表来训练模型时，每个epoch大约需要5个小时，仅使用一块 NVIDIA GeForce RTX 3070 GPU。

以下是代码：

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

八、使用训练好的模型来翻译一个日语句子

首先，我们创建用于翻译新句子的函数，其中包括获取日语句子、分词、转换为张量、推理，然后将结果解码回一个英文句子，但这次是用英文。

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)

' ▁H S ▁ 代 码 ▁85 15 ▁ 焊 接 设 备 ( 包 括 电 气 加 热 ) 。 '

trainen.pop(5)

'Chinese HS Code Harmonized Code System < HS编码 8515 : 电气(包括电热气体)、激光、其他光、光子束、超声波、电子束、磁脉冲或等离子弧焊接机器及装置,不论是否 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...'

trainja.pop(5)

'Japanese HS Code Harmonized Code System < HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)、レーザーその他の光子ビーム式、超音波式、電子ビーム式、 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...'

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

最后，在训练结束后，我们将使用Pickle先保存词汇表对象（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')

第二个也是用于推理，但还用于当我们想要稍后加载模型并继续训练的时候。

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