使用Transformer实现机器翻译

在这篇博客中，我们将探讨如何使用PyTorch实现一个基于Transformer模型的中日翻译系统。我们将使用一个双语数据集，其中包含中文和日文的句子对。这个教程涵盖数据预处理、模型定义、训练过程以及模型评估和推断。

准备工作

首先，我们需要导入必要的库，并确保设备是否可用。如果有GPU，代码会自动使用GPU进行计算。同时确保在我们的系统中安装了以下软件包，如果发现有些软件包丢失，请务必安装它们。

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

数据加载与预处理

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

我们使用pandas读取一个包含中日句子对的文本文件，并将数据转换为列表形式。

df = pd.read_csv('zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
trainzh = df[2].values.tolist()  # 中文句子
trainja = df[3].values.tolist()  # 日文句子

接下来，我们使用SentencePiece进行分词。

与英语或其他按字母顺序排列的语言不同，日语句子不包含用于分隔单词的空格。我们可以使用 JParaCrawl 提供的分词器，它是使用 SentencePiece 为日语和英语创建的，你可以访问 JParaCrawl 网站下载它们。

zh_tokenizer = spm.SentencePieceProcessor(model_file='spm.zh.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='spm.ja.nopretok.model')

我们还需要构建词汇表。

使用标记器和原始句子，构建从 TorchText 导入的词汇对象。这个过程可能需要几秒钟或几分钟，这取决于我们的数据集的大小和计算能力。

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)
zh_vocab = build_vocab(trainzh, zh_tokenizer)

将数据转换为张量，并为每个句子添加起始和结束标记。

def data_process(ja, zh):
  data = []
  for (raw_ja, raw_zh) in zip(ja, zh):
    ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    zh_tensor_ = torch.tensor([zh_vocab[token] for token in zh_tokenizer.encode(raw_zh.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    data.append((ja_tensor_, zh_tensor_))
  return data

train_data = data_process(trainja, trainzh)

定义生成批次数据的函数。

BATCH_SIZE数值的设置取决于各种事情，如您的机器内存容量，数据大小等，所以随时根据您的需要更改批量大小(注意: 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, zh_batch = [], []
  for (ja_item, zh_item) in data_batch:
    ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
    zh_batch.append(torch.cat([torch.tensor([BOS_IDX]), zh_item, torch.tensor([EOS_IDX])], dim=0))
  ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
  zh_batch = pad_sequence(zh_batch, padding_value=PAD_IDX)
  return ja_batch, zh_batch

train_iter = DataLoader(train_data, batch_size=BATCH_SIZE, shuffle=True, collate_fn=generate_batch)

定义Transformer模型

接下来，我们定义一个基于Transformer的Seq2Seq模型。

其中，变压器是一个 Seq2Seq 模型介绍了“注意力是你所需要的一切”文件，以解决机器翻译任务。变压器模型由一个编码器和解码器块组成，每个编码器和解码器块包含固定数量的层。

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

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

模型训练与评估

定义训练和评估函数。

最后，在准备好必要的类和函数之后，我们就可以开始训练模型了。这是不言而喻的，但是完成训练所需的时间可能会因计算能力、参数和数据集大小等许多因素而有很大的不同。

当使用来自 JParaCrawl 的完整句子列表来训练这个模型时，每种语言大约有590万个句子，使用一个单独的 NVIDIA GeForce RTX 3070图形处理器大约需要5个小时。

SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(zh_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)

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)

      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)

for epoch in 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:.4f}, Epoch time = {(end_time - start_time):.4f}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)
        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).fill_(next_word).type(torch.long).to(device)], dim=0)
        if next_word == EOS_IDX:
            break
    return ys

def translate(model: torch.nn.Module, src_sentence: str):
    model.eval()
    src = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(src_sentence, out_type=str)],
                       dtype=torch.long).reshape(-1, 1)
    num_tokens = src.shape[0]
    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(zh_vocab.lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace("<bos>", "").replace("<eos>", "")

现在，我们可以使用训练好的模型来生成翻译了

translate(transformer, "こんにちは")