基于Transformer实现的机器翻译（日译中）

1 实验内容简介

实现一个简单的基于Transformer机构的日译中翻译模型。

2 Transformer架构

Transformer是一种基于自注意力机制的神经网络架构，由编码器和解码器两部分组成，广泛应用于自然语言处理任务。其核心思想是通过自注意力机制（Self-Attention）计算输入序列中每个元素对其他元素的相关性，从而捕捉长距离依赖关系。为了保留序列中的位置信息，Transformer引入了位置编码（Positional Encoding）。此外，多头注意力机制（Multi-Head Attention）通过并行计算多个注意力头，使模型能够关注不同的子空间，提高了模型的表达能力和效果。

2.1 自注意力机制（self-attention）

首先，通过一个例子从宏观角度理解自注意力机制的工作原理与重要意义。
假设要将“His hair was so long that he was going to the barbershop in the afternoon to cut it off.”翻译为中文。作为人类，我们一眼就可以看出句子尾端的it指的是his hair。但是在计算机的视角里，它只是一串二进制代码，计算机很难知道二者之间的联系。再者，对于一词多义现象，我们很容易地可以根据上下文信息准确地判断该词到底是什么意思，但是这对于计算机来说，同样是极具挑战性的任务。

2.1.1 自注意力机制的工作原理

那么自注意力机制该如何处理这些问题呢？简单来说，就是通过计算该词与句子中其他词语的相似度，并根据相似度赋予对应的权重。相似度越高则权重越高，也就意味着词语之间相关度更高。然后再将这些词语的向量表示进行加权求和，以此将句子中所有词的信息按照相关程度揉合为该词的向量表示。
计算自注意力有两种方式：一种通过向量，一种通过矩阵。二者原理相同，通过矩阵运算只是将通过向量计算的所有向量组成矩阵，使用矩阵运算方法进行计算而已。

2.1.2自注意力机制的计算步骤

1. 对输入向量进行线性变换：
   对于每一个输入向量$x_i$，首先通过三个不同的线性变换得到查询向量$q_i$、键向量$k_i$和值向量$v_i$：
   $$Q=X{W}^Q，K=X{W}^K，V=X{W}^V$$

• 其中$X$是输入序列的矩阵表示，$W^Q$、$W^K$和$W^V$是可学习的权重矩阵。
• 查询向量Query是当前单词的表示形式，用于对所有其他单词(key)进行评分。
• 键向量Key可以看作是序列中所有单词的标签，是我们在找相关单词时候的对照物。
• 值向量Value是单词的实际表示，一旦我们对每个单词的相关度打分之后，我们就要对Value进行相加表示当前正在处理的单词的value。

2. 计算注意力得分：
   查询向量$q_i$和所有键向量$k_i$之间的相似度使用点积计算，并进行缩放：
   $$Attention(Q,K,V)=softmax(\frac{Q{K}^T}{\sqrt{d_k}})V$$
   其中，$\sqrt{d_k}是缩放因子，$$d_k$是键向量的维度，目的是缓解点积值过大的问题。
3. 加权求和得到输出向量：
   $$Attention(Q,K,V)=\sum _{j=1}^{n}softmax(\frac{q_i\cdot k_j}{\sqrt{d_k}})v_j$$

2.2 多头注意力机制（multi-head attention）

在自注意力机制的基础上，Transformer引入了多头注意力机制（Multi-Head Attention）。多头注意力机制将查询、键和值向量分成多个头（head），每个头独立计算注意力，然后将结果拼接起来，再通过线性变换得到最终的输出。以探索不同的子空间，捕捉序列中更丰富的特征信息。
$$MultiHead(Q,K,V)=Concat(hea{d}_1,hea{d}_2,...,hea{d}_h)W^O$$
其中每个头的计算为：
$$hea{d}_i=Attention(Q{W}_i^Q,K{W}_i^K,V{W}_i^V)$$
$W_i^Q$,$W_i^K$,$W_i^V$和$W_i^O$是可学习的权重矩阵。

2.3 位置编码（positional encoding）

由上述介绍不难看出，在计算自注意力得分时，输入序列中每个位置的单词都各自单独的路径流入编码器，即各个单词同时流入编码器中，不是排队进入。而这严重影响着句子的意思。比如，“小明喜欢我”和“我喜欢小明”，这两个句子只是，词序不同，却表达着截然相反的意思。此时再结合RNN的结构不难发现Transformer完全把时序信息给丢掉了。
为了解决时序的问题，Transformer的作者用了一个绝妙的办法：位置编码（Positional Encoding）。
即给每个位置编号，从而每个编号对应一个向量，最终通过结合位置向量和词向量，作为输入embedding，就给每个词都引入了一定的位置信息，这样Attention就可以分辨出不同位置的词了。

3 实验步骤

3.1导入必要的包

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)) ## 如果你有GPU，请在你自己的电脑上尝试运行这一套代码

3.2 获取并行数据集

在本实验中，我们将使用从 JParaCrawl 下载的日英并行数据集！[http://www.kecl.ntt.co.jp/icl/lirg/jparacrawl]，它是NTT创建的最大的公开可用的英日平行语料库。

df = pd.read_csv('./zh-ja/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)
#中文句子列表
trainchn = df[2].values.tolist()#[:10000]
#英文句子列表
trainja = df[3].values.tolist()#[:10000]
# trainen.pop(5972)
# trainja.pop(5972)
print(trainchn[500])
print(trainja[500])

运行结果：

3.3准备分词器

与英语不同，日语句子不包含空格来分隔单词。我们可以使用JParaCrawl提供的分词器，该分词器使用 SentencePiece 分别为日语和英语创建。

chn_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')
# chn_tokenizer.encode("All residents aged 20 to 59 years who live in Japan must enroll in public pension system.", out_type='str')
chn_tokenizer.encode("自然语言处理好难！",out_type='str')
# chn_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')

运行结果：

3.4 构建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)
chn_vocab = build_vocab(trainchn, chn_tokenizer)

在我们成功构建词汇表之后，就可以使用词汇表（vocab）和分词器（tokenizer）来为训练数据构建张量（tensors）了。我们需要将文本数据转换成数值形式，以便神经网络能够处理。具体来说，这个过程通常包括将句子切分成单词或子词，将每个单词或子词映射到词汇表中的一个唯一索引，然后将这些索引转换成可以输入到模型中的张量格式。这样，文本数据就被表示成了可以被机器学习模型理解的数值型数据结构。

#数据预处理
def data_process(ja, chn):
  data = []
  #遍历每一对中文与英文句子
  for (raw_ja, raw_chn) in zip(ja, chn):
    #转换为tensor张量
    ja_tensor_ = torch.tensor([ja_vocab[token] for token in ja_tokenizer.encode(raw_ja.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    chn_tensor_ = torch.tensor([chn_vocab[token] for token in chn_tokenizer.encode(raw_chn.rstrip("\n"), out_type=str)],
                            dtype=torch.long)
    data.append((ja_tensor_, chn_tensor_))
  return data
train_data = data_process(trainja, trainchn)

3.5 构建DataLoader以生成训练数据

BATCH_SIZE = 8   #批处理大小
#获取词汇表中<pad><bos><eos>的索引
PAD_IDX = ja_vocab['<pad>']
BOS_IDX = ja_vocab['<bos>']
EOS_IDX = ja_vocab['<eos>']

#生成一个批次的训练数据
def generate_batch(data_batch):
  ja_batch, chn_batch = [], []
  for (ja_item, chn_item) in data_batch:
    #在每个句子前后添加<bos><eos>标记
    ja_batch.append(torch.cat([torch.tensor([BOS_IDX]), ja_item, torch.tensor([EOS_IDX])], dim=0))
    chn_batch.append(torch.cat([torch.tensor([BOS_IDX]), chn_item, torch.tensor([EOS_IDX])], dim=0))
  #填充<pad>,对齐批次中的每个句子
  ja_batch = pad_sequence(ja_batch, padding_value=PAD_IDX)
  chn_batch = pad_sequence(chn_batch, padding_value=PAD_IDX)
  return ja_batch, chn_batch
#使用DataLoader加载训练数据
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

3.6 序列到序列（Sequence-to-Sequence）变换器

接下来的几段代码和文本解释（用斜体书写）摘自原始的 PyTorch 教程 [https://pytorch.org/tutorials/beginner/translation_transformer.html]。除了将 BATCH_SIZE 和词汇表 de_vocab 更改为 ja_vocab 外，我没有进行任何更改。

变换器（Transformer）是一种序列到序列（Seq2Seq）模型，最初在论文“Attention is all you need”中引入，用于解决机器翻译任务。Transformer 模型包括编码器（encoder）和解码器（decoder）模块，每个模块包含固定数量的层。

编码器通过一系列多头注意力（Multi-head Attention）和前馈网络层处理输入序列。编码器的输出称为记忆（memory），将其与目标张量一起馈送给解码器。编码器和解码器通过教师强制（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__()
        #初始化编码器层
        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 embeddings）表示。为了恢复句子中单词序列的顺序，位置编码（positional encoding）被添加到标记嵌入中。

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)

然后创建一个后续词屏蔽（subsequent word mask），用于阻止目标词语注意其后续的词语。同时还创建了用于屏蔽源语言和目标语言填充标记的屏蔽（mask）。

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

设置超参数，实例化模型，构建模型训练与评估参数。

#超参数
SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(chn_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)
#Adam优化算法
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)

3.7开始训练

最后，在准备了所有必需的类和函数后，我们就准备好训练我们的模型了。受多种因素影响，完成训练所需的时间会有很大差异，比如计算能力、模型参数以及数据集的大小。优化这些因素不仅可以缩短训练时间，还能影响模型的性能和准确性。因此，在开始训练之前，合理配置资源、调整超参数以及对数据进行有效的预处理都是至关重要的步骤。训练过程中，我们通常会监控损失函数（loss function）的变化以及验证集上的性能，以评估模型的学习进度并适时作出调整。
首先，这里给出我在进行模型训练时的电脑配置与版本信息。

def print_device_info():
    # 打印 PyTorch 相关库的版本
    print(f"PyTorch version: {torch.__version__}")
    print(f"torchvision version: {torchvision.__version__}")
    print(f"torchaudio version: {torchaudio.__version__}")
    print(f"torchtext version: {torchtext.__version__}")
    print(f"sentencepiece version: {sentencepiece.__version__}")

    # 检查并打印 CUDA 相关信息
    if torch.cuda.is_available():
        print(f"CUDA is available: {torch.cuda.is_available()}")
        print(f"CUDA version: {torch.version.cuda}")
        print(f"CUDNN version: {torch.backends.cudnn.version()}")
        print(f"Number of GPUs: {torch.cuda.device_count()}")
        for i in range(torch.cuda.device_count()):
            print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
            print(f"GPU {i} Capability: {torch.cuda.get_device_capability(i)}")
    else:
        print("CUDA is not available.")

    # 打印设备信息
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    print(f"Device: {device}")

运行结果：

开始训练

# 开始训练
print("strat training!")
for epoch in range(1, NUM_EPOCHS + 1):
    # 进行一个epoch的训练，并记录训练损失
    train_loss = train_epoch(transformer, train_iter, optimizer)
    end_time = time.time()
    # 输出当前epoch的训练信息
    print((f"Epoch:{epoch},Train loss:{train_loss:.3f},"
           f"Epoch time={(end_time - start_time):.3f}s"))

    # 保存当前 epoch 的模型信息
    torch.save({
        'epoch': epoch,
        'model_state_dict': transformer.state_dict(),
        'optimizer_state_dict': optimizer.state_dict(),
        'loss': train_loss,
    }, f'models2/epoch_{epoch}_model_checkpoint.tar')

print("Finish!")
print(f"Training started at: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(start_time))}")
end_time = time.time()
print(f"Training ended at: {time.strftime('%Y-%m-%d %H:%M:%S', time.localtime(end_time))}")

训练过程截图：

收个人笔记本配置限制，在训练时batch_size只能取到4。因此每一个epoch都需要大概20分钟的时间。下面时训练结束后保存的学得模型：

3.8 尝试使用学得模型翻译日文句子

首先，我们需要创建一些函数来翻译新句子，这包括以下几个步骤：获取日语句子、进行分词、转换为张量、进行推理预测，然后将结果解码回句子形式。

#实现贪婪解码，从模型生成翻译结果
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)

翻译结果：

That’s all.