基于Transformer的简单机器翻译

WA2114084, AHU UNIVERSITY

前言

Transformer是一种用于自然语言处理和其他序列到序列任务的神经网络模型。它引入了自注意力机制来实现序列之间的建模，取代了传统的循环神经网络。

Transformer已经在机器翻译、文本摘要、语言模型等任务中取得了很好的效果，并且被广泛应用和研究。它的引入和成功应用为自然语言处理领域带来了重大的突破和进展。

本文将基于Transformer模型，构建一个用于中-日翻译的机器翻译模型。

 

准备环境

首先，我们需要准备好运行我们的模型所需的环境。

硬件

你最好有一块至少有8GB VRAM的，支持CUDA的英伟达GPU。如果能支持半精度计算就更好了。RTX系列的GPU（包括游戏卡和专业卡）差不多都支持混合精度，使用这玩意能大幅降低显存占用，加快训练速度。如果没有GPU，你可以考虑去诸如 AutoDL 之类的算力平台上租用GPU。一台含4090(24G VRAM)的机器每小时的价格大概在1.88￥左右。

软件

你需要确保安装了python，torchtext 和 sentencepiece。

torchtext和sentencepiece是我们这次训练模型所必须的两个包。其他包会在这两个包被安装时，作为依赖同时被安装。

我使用的是python==3.10.0， 确保torchtext==0.6.0, sentencepiece==0.2.0即可。使用对应的版本可以避免API不兼容的麻烦。

pip install torchtext==0.6.0 sentencepiece==0.2.0 pandas

如果你使用了jupyter notebook,也请确保安装了必要的依赖。

接下来，导入我们所需要的所有包吧。如果没有报错，就说明你的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")
print(torch.cuda.get_device_name(0))

准备数据集

我们的数据集会从JParaCrawl (http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl) 这里搞到。

当然,如果你不想下载未清理的原始数据集的话，这里有现成的（33MB）数据集。

请下载这个文件并解压到你的工作目录里，里面包含数据集，以及基础模型。

中文-日文_机器翻译模型与数据集.7z 密码： 23jx文件大小：7.8 M|https://wwb.lanzouq.com/iF43G232qzaj密码:23jx

（蓝奏云没办法关闭下载密码，CSDN则有概率把你的资源变成会员付费专享。对此造成的麻烦实在抱歉）

（另外还请不要细看数据集里面的内容，不太干净，多看一眼就会爆炸 ... ）

 

处理词汇

首先,让我们看看我们的数据吧。

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]

print(trainen[114], "\n", trainja[114])
print(trainen[514], "\n", trainja[514])

输出

Chinese HS Code Harmonized Code System < HS编码 2914 酮及醌及其卤化、磺化、硝化或亚硝化衍生物 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コード 2914 ケトン及びキノン(他の酸素官能基を有するか有しないかを問わない。)並びにこれらのハロゲン化誘導体、スルホン化誘導体、ニトロ化誘導体 HS Code List (Harmonized System Code) for US, UK, EU, China, India, France, Japan, Russia, Germany, Korea, Canada ...
FARO宣布推出适用于BIM-CIM和公共安全领域的新款HDR激光扫描仪和即时渲染软件 | FARO 法如中国 佛罗里达州玛丽湖,2016年4月11日 – 在BIM-CIM 和公共安全市场领域, 最受全球用户信赖的三维测量、成像和实现技术供应商FARO Technologies, Inc.(NASDAQ:FARO)宣布推出具有高动态范围成像(HDR)功能的新型三维激光扫描仪,这款激光扫描仪能够在明亮或黑暗的光照条件下生成无比清晰的图像,并具有更高的分辨率和更佳的色彩显示效果。 
 FARO、BIM/CIMおよび公共の安全のための用途に最適な、新しいHDR レーザースキャナーと、インスタントレンダリングソフトウェアを発表 | FARO ファロージャパン フロリダ州レイク・メアリー、2016年4月11日―BIM/CIMおよび公共の安全のための用途において、世界で最も信頼のおける3次元測定技術とイメージングおよびリアル化技術を提供するFARO Technologies, Inc.

准备Tokenizer

不同于英语，日语句子不包含空格来分隔单词（和中文一样）。我们可以使用JParaCrawl提供的分词器，该分词器使用SentencePiece为日语和英语分别创建，你可以访问JParaCrawl网站下载它们，或者用我在“准备数据集”里提供的那个压缩包内的模型，我已经事先为你准备好了。

只需要解压我上文中提供的压缩包到你的工程根目录即可。你需要 spm.ja.nopretok.model 和spm.en.nopretok.model 这两个文件。确保它们在你的工程根目录里。

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

加载完成后，你可以测试这俩分词器的效果。比如这样：

print(en_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type=str))
print(ja_tokenizer.encode("年金 日本に住んでいる20歳~60歳の全ての人は、公的年金制度に加入しなければなりません。", out_type=str))

输出：

['▁All', '▁residents', '▁aged', '▁20', '▁to', '▁59', '▁years', '▁who', '▁live', '▁in', '▁Japan', '▁must', '▁enroll', '▁in', '▁public', '▁pension', '▁system', '.']
['▁', '年', '金', '▁日本', 'に住んでいる', '20', '歳', '~', '60', '歳の', '全ての', '人は', '、', '公的', '年', '金', '制度', 'に', '加入', 'しなければなりません', '。']

构建TorchText Vocab对象

使用分词器和原始句子，构建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)
en_vocab = build_vocab(trainen, en_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("\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.append((ja_tensor_, en_tensor_))
  return data
train_data = data_process(trainja, trainen)

准备数据加载器

为了防止爆显存，还请按需修改BATCH_SIZE。

开启混合精度可以大幅减少显存消耗，这样你就可以把BATCH_SIZE开得更高了。这点我们会在后面提到。

BATCH_SIZE = 32
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
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

准备我们的Transformer模型

接下来的几个代码和解释取自原始的PyTorch教程，几乎没有做任何更改：

Language Translation with nn.Transformer and torchtext — PyTorch Tutorials 2.3.0+cu121 documentationhttps://pytorch.org/tutorials/beginner/translation_transformer.html

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

关于Transformer的详细的介绍可以参考互联网上的教程，比我这里的要详细的多。

以下代码差不多原封不动取自Torch的官方教程：

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)

文字Token使用embedding来表示。由于结构特殊，在transformer里，我们需要单独加入位置编码来表示单词的顺序。

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

训练模型，启用混合精度

定义模型，定义超参数，设置优化器和损失函数，启用混合精度。

Torch提供了一个方便的办法来让我们实现混合精度。你可以简单得使用autocast来创建混合精度的上下文环境，用GradScaler来对精度进行缩放，即可实现混合精度的计算。

SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(en_vocab)
EMB_SIZE = 512
NHEAD = 8
FFN_HID_DIM = 512
# BATCH_SIZE = 16
BATCH_SIZE = 64
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
)


from torch.cuda.amp import GradScaler, autocast
scaler = GradScaler()  # 初始化GradScaler

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)

        # 使用autocast上下文管理器自动处理数据类型转换
        with autocast():
            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))

        # 使用GradScaler来缩放损失，进行反向传播
        optimizer.zero_grad()
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        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)

        # 使用autocast上下文管理器自动处理数据类型转换
        with autocast():
            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)

来开始训练吧。

在我的（租的）RTX4090上，训练完毕该模型花了我不到20分钟。每个epoch大约耗时65秒。

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 ▁ 焊 接 或 焊 接 设 备 ( 包 括 电 气 加 热 器 ) 。 '

好像有那么回事，但总觉得有点驴头不对马嘴... 这大概是分词器不太对？还请各位读者自行寻找改良的方法吧，没精力优化了

保存结果

首先保存Vocab对象（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')

保存完毕的推理模型大概占据了178MB的空间，检查点模型则占据了515MB空间。

工程文件

如果你懒得阅读全文，或者想直接运行该项目而不是一段一段复制，这里有已经配置完毕的工程项目：

基于Transformer的中-日机器翻译-工程.zip - 蓝奏云文件大小：13.9 M|https://wwb.lanzouq.com/i8V43232so6b