使用pytorch和transformer完成日译中的翻译任务

目录

     1.实验介绍

     2.实验内容

2.1获取平行数据集

2.2准备分词器

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

2.4开始训练

2.4保存vocab对象和训练模型

     3.实验总结

1.实验介绍

随着全球化的不断深入，语言障碍成为了跨文化交流的一大难题。机器翻译作为解决这一问题的关键技术，近年来得到了飞速的发展。特别是Transformer模型的提出，以其自注意力机制在处理序列数据方面展现出了卓越的能力，已成为NLP领域的新宠。PyTorch是一个广泛使用的开源机器学习库，以其动态计算图和易用性而受到研究者和开发者的青睐。而Transformer模型，自从在2017年被Google提出后，已经在机器翻译等多个NLP任务中取得了突破性进展。

2.实验内容

导入相关库

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

2.1获取平行数据集

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

df = pd.read_csv('./zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

trainen = df[2].values.tolist()#[:10000]

trainja = df[3].values.tolist()#[:10000]

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

这里是数据集中包含的一条句子的例子。

print(trainen[500])

print(trainja[500])

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

2.2准备分词器

与英语或其他字母语言不同，日语句子中没有空格来分隔单词。我们可以使用由JParaCrawl提供的分词器，它是使用SentencePiece创建的，适用于日语和英语。您可以前往JParaCrawl网站下载它们。

en_tokenizer = spm.SentencePieceProcessor(model_file='./spm.en.nopretok.model')

ja_tokenizer = spm.SentencePieceProcessor(model_file='./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')

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

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

def build_vocab(sentences, tokenizer):

  counter = Counter()

  for sentence in sentences:

    counter.update(tokenizer.encode(sentence, out_type=str))#对于每个句子，使用 tokenizer.encode 方法进行编码，将句子转换为词汇的序列。out_type=str 参数指定输出类型为字符串。然后使用 counter.update 方法更新计数器，将编码后的句子中的词汇加入计数。

  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("\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)#对于每一对日语和英语句子，使用分词器 ja_tokenizer 和 en_tokenizer 将句子分词并编码为词汇索引序列。raw_ja.rstrip("\n") 和 raw_en.rstrip("\n") 用于去除句子末尾的换行符。ja_tokenizer.encode 和 en_tokenizer.encode 方法将句子转换为词汇索引序列。使用列表推导式 [ja_vocab[token] for token in ...] 和 [en_vocab[token] for token in ...] 将词汇索引转换为对应的词汇表索引。

    data.append((ja_tensor_, en_tensor_))

  return data

train_data = data_process(trainja, trainen)

在训练过程中创建DataLoader对象进行迭代。

在这里，我将BATCH_SIZE设置为16以防止“cuda内存不足”，但这取决于诸如您的机器内存容量、数据大小等各种因素，因此根据您的需求随时可以更改批量大小（注：PyTorch的教程将批量大小设置为128，使用了Multi30k德英数据集作为示例）。

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:#循环处理 data_batch 中的每个元素 (ja_item, en_item):对于每对日语和英语句子的张量，使用 torch.cat 将 <bos> 和 <eos> 标记添加到句子的开始和结束位置。将处理后的张量添加到 ja_batch 和 en_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)#使用 pad_sequence 函数对 ja_batch 和 en_batch 中的句子进行填充，使它们具有相同的长度。padding_value=PAD_IDX 指定使用 <pad> 标记进行填充。

  return ja_batch, en_batch

train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,

                        shuffle=True, collate_fn=generate_batch)

Sequence-to-sequence Transformer

下面几段代码和文本解释（以斜体字写）取自原始的PyTorch教程[https://pytorch.org/tutorials/beginner/translation_transformer.html]。我没有进行任何更改，只修改了BATCH_SIZE和单词de_vocab改为ja_vocab。

Transformer是一种Seq2Seq模型，被引入到“Attention is all you need”论文中，用于解决机器翻译任务。Transformer模型由一个包含固定层数的编码器块和一个包含固定层数的解码器块组成。

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

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) # 解码器处理目标序列，并考虑编码器的输出（memory） 

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

         # 使用torch.arange创建等差数列，并计算分母

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

         # 使用sin和cos函数计算位置嵌入

        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)

    # 将矩阵的上三角部分设置为True（1），下三角部分设置为False（0），然后转置行列，使下三角为True 

    # 将布尔值转换为浮点数，上三角（False部分）填充为负无穷（-inf），下三角（True部分）填充为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)

  # 源序列之间不需要进行掩码（假设源序列的所有标记都是可见的），因此初始化为全0的布尔矩阵   

  src_mask = torch.zeros((src_seq_len, src_seq_len), device=device).type(torch.bool)

  # 生成源序列的填充掩码，标记PAD_IDX位置为True（即需要被忽略的位置）

  src_padding_mask = (src == PAD_IDX).transpose(0, 1)

  # 生成目标序列的填充掩码，同样标记PAD_IDX位置为True 

  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 = 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)# 使用Xavier初始化对模型参数进行初始化 

transformer = transformer.to(device)

loss_fn = torch.nn.CrossEntropyLoss(ignore_index=PAD_IDX)# 定义损失函数，忽略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)

2.4开始训练

最终，在准备好必要的类和函数之后，我们已准备好训练我们的模型了。毋庸置疑，完成训练所需的时间会因计算能力、参数和数据集大小等因素而有很大不同。

当我使用来自JParaCrawl的完整句子列表进行模型训练时，每种语言约有590万个句子，仅使用一块 NVIDIA GeForce RTX 3070 GPU，每个epoch 大约需要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"))

请尝试使用训练模型翻译一句日语句子。首先，我们创建翻译新句子的功能，包括获取日语句子、分词、转换为张量、推理，然后将结果解码回句子，但这次是用英语。

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)

2.5保存vocab对象和训练模型

在训练结束后，我们将首先使用Pickle保存词汇表对象（en_vocab和ja_vocab）

import pickle

# open a file, where you want to store the data

file = open('en_vocab.pkl', 'wb')

# dump information to that file

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

3.实验总结

Transformer模型是一种基于自注意力机制的架构，它摒弃了传统的循环神经网络（RNN）结构，能够并行处理序列数据，从而显著提高了训练效率，尤其是在使用现代GPU加速时。

使用Transformer的模型通常能够提供更好的翻译性能，特别是在处理长距离依赖关系时。自注意力机制使得模型能够捕捉到输入序列中任意两个位置之间的依赖关系，无论它们之间的距离有多远。