基于神经网络的机器翻译——以Transformer为例

Transformer是一种基于注意力机制（attention mechanism）的深度学习模型，最初被用于自然语言处理任务，特别是机器翻译。Transformer的主要特点在于自注意力机制（self-attention），它能够同时计算序列中所有位置的注意力权重，从而更好地捕捉长距离依赖关系。

注意力机制

循环神经网络和卷积神经网络一半用来处理表格或者图像形式的数据，默认数据都来自于某种分布，并且所有样本都是独立同分布的（independently and identically distributed，i.i.d.）。然而，大多数的数据并非如此。例如，文章中的单词是按顺序写的，如果顺序被随机地重排，就很难理解文章原始的意思。同样，视频中的图像帧、对话中的音频信号以及网站上的浏览行为都是有顺序的。因此，针对此类数据而设计特定模型，可能效果会更好。因此，在这个基础上我们引入注意力机制人为确定哪些数据是相对重要的。

注意力机制通过注意力汇聚（attentionpooling）将选择引导至感官输入（sensoryinputs，例如中间特征表示）。在注意力机制中，这些感官输入被称为值（value）。更通俗的解释，每个值都与一个键（key）配对，这可以想象为感官输入的非自主提示。如图所示，可以通过设计注意力汇聚的方式，便于给定的查询（自主性提示）与键（非自主性提示）进行匹配，这将引导得出最匹配的值（感官输入）。

有了注意力机制之后，我们将词元序列输入注意力池化中，以便同一组词元同时充当查询、键和值。具体来说，每个查询都会关注所有的键值对并生成一个注意力输出。由于查询、键和值来自同一组输入，因此被称为自注意力，也被称为内部注意力。

3. 机器翻译构建实例

3.1 读取和预处理数据

定义一些特殊符号。其中“<pad>”（padding）符号用来添加在较短序列后，直到每个序列等长，而“<bos>”和“<eos>”符号分别表示序列的开始和结束。

import collections  # 导入collections模块，用于高效的数据容器
import os  # 导入os模块，用于与操作系统进行交互
import io  # 导入io模块，用于处理流数据
import math  # 导入math模块，提供数学运算功能
import torch  # 导入PyTorch深度学习框架
from torch import nn  # 从torch中导入nn模块，用于构建神经网络
import torch.nn.functional as F  # 导入torch.nn.functional模块，提供神经网络函数
import torchtext.vocab as Vocab  # 导入torchtext.vocab模块，用于处理文本词汇
import torch.utils.data as Data  # 导入torch.utils.data模块，提供数据加载和处理工具

import sys  # 导入sys模块，用于系统相关的参数和函数
# sys.path.append("..") 
import d2lzh_pytorch as d2l  # 导入d2lzh_pytorch模块，引入相关深度学习工具函数

PAD, BOS, EOS = '<pad>', '<bos>', '<eos>'  # 定义特殊词汇常量，用于标记填充、序列开始和序列结束
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # 设置环境变量CUDA_VISIBLE_DEVICES为0，指定使用GPU编号为0
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 根据是否有CUDA设备选择使用GPU或CPU

print(torch.__version__, device)  # 打印当前PyTorch版本号和设备信息

接着定义两个辅助函数对后面读取的数据进行预处理。

# 将一个序列中所有的词记录在all_tokens中以便之后构造词典，
# 然后在该序列后面添加PAD直到序列长度变为max_seq_len，
# 然后将序列保存在all_seqs中
def process_one_seq(seq_tokens, all_tokens, all_seqs, max_seq_len):
    all_tokens.extend(seq_tokens)  # 将当前序列的词加入到总词汇表中
    seq_tokens += [EOS] + [PAD] * (max_seq_len - len(seq_tokens) - 1)  # 在序列末尾添加EOS和PAD直到序列达到最大长度
    all_seqs.append(seq_tokens)  # 将处理后的序列加入到总序列列表中

# 使用所有的词来构造词典，并将所有序列中的词转换为词索引后构造Tensor
def build_data(all_tokens, all_seqs):
    vocab = Vocab.Vocab(collections.Counter(all_tokens), specials=[PAD, BOS, EOS])  # 使用所有词构建词典，包括特殊词汇PAD、BOS、EOS
    indices = [[vocab.stoi[w] for w in seq] for seq in all_seqs]  # 将所有序列中的词转换为词索引
    return vocab, torch.tensor(indices)  # 返回词典对象和包含词索引的Tensor

def read_data(max_seq_len):
    # in和out分别是input和output的缩写
    in_tokens, out_tokens, in_seqs, out_seqs = [], [], [], []
    with io.open('fr-en-small.txt') as f:
        lines = f.readlines()  # 读取文件中的所有行
    for line in lines:
        in_seq, out_seq = line.rstrip().split('\t')  # 按照制表符分割输入和输出序列
        in_seq_tokens, out_seq_tokens = in_seq.split(' '), out_seq.split(' ')  # 分割输入和输出序列中的词汇
        if max(len(in_seq_tokens), len(out_seq_tokens)) > max_seq_len - 1:  # 如果加上EOS后长于max_seq_len，则忽略掉此样本
            continue
        process_one_seq(in_seq_tokens, in_tokens, in_seqs, max_seq_len)  # 处理输入序列
        process_one_seq(out_seq_tokens, out_tokens, out_seqs, max_seq_len)  # 处理输出序列
    in_vocab, in_data = build_data(in_tokens, in_seqs)  # 构建输入词典和Tensor
    out_vocab, out_data = build_data(out_tokens, out_seqs)  # 构建输出词典和Tensor
    return in_vocab, out_vocab, Data.TensorDataset(in_data, out_data)  # 返回输入和输出的词典以及包含Tensor的数据集对象

将序列的最大长度设成7，然后查看读取到的第一个样本。该样本分别包含法语词索引序列和英语词索引序列。

max_seq_len = 7  # 设置最大序列长度为7

# 调用read_data函数读取数据集，返回输入词汇表、输出词汇表和数据集
in_vocab, out_vocab, dataset = read_data(max_seq_len)

# 输出数据集中的第一个样本
print(dataset[0])

3.2 含注意力机制的编码器—解码器

我们将使用含注意力机制的编码器—解码器来将一段简短的法语翻译成英语。下面我们来介绍模型的实现。

3.2.1 编码器

在编码器中，我们将输入语言的词索引通过词嵌入层得到词的表征，然后输入到一个多层门控循环单元中。PyTorch的nn.GRU实例在前向计算后也会分别返回输出和最终时间步的多层隐藏状态。其中的输出指的是最后一层的隐藏层在各个时间步的隐藏状态，并不涉及输出层计算。注意力机制将这些输出作为键项和值项。

class Encoder(nn.Module):
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 drop_prob=0, **kwargs):
        super(Encoder, self).__init__(**kwargs)
        self.embedding = nn.Embedding(vocab_size, embed_size)  # 定义词嵌入层，将词汇索引映射为稠密词向量
        self.rnn = nn.GRU(embed_size, num_hiddens, num_layers, dropout=drop_prob)  # 定义GRU循环神经网络层

    def forward(self, inputs, state):
        # 输入形状是(批量大小, 时间步数)。将输出互换样本维和时间步维
        embedding = self.embedding(inputs.long()).permute(1, 0, 2)  # 执行词嵌入操作，并调整张量维度
        return self.rnn(embedding, state)  # 返回GRU层的输出和新的状态

    def begin_state(self):
        return None  # 初始状态返回None，适用于单向GRU的情况

下面我们来创建一个批量大小为4、时间步数为7的小批量序列输入。设门控循环单元的隐藏层个数为2，隐藏单元个数为16。编码器对该输入执行前向计算后返回的输出形状为(时间步数, 批量大小, 隐藏单元个数)。门控循环单元在最终时间步的多层隐藏状态的形状为(隐藏层个数, 批量大小, 隐藏单元个数)。对于门控循环单元来说，state就是一个元素，即隐藏状态；如果使用长短期记忆，state是一个元组，包含两个元素即隐藏状态和记忆细胞。

encoder = Encoder(vocab_size=10, embed_size=8, num_hiddens=16, num_layers=2)

# 执行编码器的前向传播，输入是一个全零张量，形状为(4, 7)，初始状态为None
output, state = encoder(torch.zeros((4, 7)), encoder.begin_state())

# 输出编码器的输出张量形状和状态张量形状
print(output.shape, state.shape)  # GRU的state是h，而LSTM的是一个元组(h, c)

3.2.2 注意力机制

我们将输入连结后通过含单隐藏层的多层感知机变换。其中隐藏层的输入是解码器的隐藏状态与编码器在所有时间步上隐藏状态的一一连结，且使用tanh函数作为激活函数。输出层的输出个数为1。两个Linear实例均不使用偏差。其中函数a𝑎定义里向量v𝑣的长度是一个超参数，即attention_size。

import torch.nn as nn

def attention_model(input_size, attention_size):
    """
    定义一个注意力模型。

    参数:
    input_size (int): 输入特征的维度。
    attention_size (int): 注意力模型的隐藏层维度。

    返回:
    model (nn.Module): 定义好的注意力模型。
    """
    # 创建一个顺序模型，即一个包含多个层（或模块）的序列。
    model = nn.Sequential(
        # 第一个全连接层，输入维度为input_size，输出维度为attention_size，不使用偏置项。
        nn.Linear(input_size, attention_size, bias=False),
        # 应用tanh激活函数。
        nn.Tanh(),
        # 第二个全连接层，输出维度为1，不使用偏置项。
        nn.Linear(attention_size, 1, bias=False)
    )
    # 返回定义好的注意力模型。
    return model

注意力机制的输入包括查询项、键项和值项。设编码器和解码器的隐藏单元个数相同。这里的查询项为解码器在上一时间步的隐藏状态，形状为(批量大小, 隐藏单元个数)；键项和值项均为编码器在所有时间步的隐藏状态，形状为(时间步数, 批量大小, 隐藏单元个数)。注意力机制返回当前时间步的背景变量，形状为(批量大小, 隐藏单元个数)。

def attention_forward(model, enc_states, dec_state):
    """
    enc_states: (时间步数, 批量大小, 隐藏单元个数)
    dec_state: (批量大小, 隐藏单元个数)
    """
    # 将解码器隐藏状态广播到和编码器隐藏状态形状相同后进行连结
    dec_states = dec_state.unsqueeze(dim=0).expand_as(enc_states)
    enc_and_dec_states = torch.cat((enc_states, dec_states), dim=2)
    e = model(enc_and_dec_states)  # 形状为(时间步数, 批量大小, 1)
    alpha = F.softmax(e, dim=0)  # 在时间步维度做softmax运算
    return (alpha * enc_states).sum(dim=0)  # 返回背景变量

在下面的例子中，编码器的时间步数为10，批量大小为4，编码器和解码器的隐藏单元个数均为8。注意力机制返回一个小批量的背景向量，每个背景向量的长度等于编码器的隐藏单元个数。因此输出的形状为(4, 8)。

seq_len, batch_size, num_hiddens = 10, 4, 8
model = attention_model(2*num_hiddens, 10)  # 创建一个注意力模型
enc_states = torch.zeros((seq_len, batch_size, num_hiddens))  # 初始化编码器状态
dec_state = torch.zeros((batch_size, num_hiddens))  # 初始化解码器状态
attention_forward(model, enc_states, dec_state).shape  # 计算注意力权重并返回形状

3.2.3 含注意力机制的解码器

我们直接将编码器在最终时间步的隐藏状态作为解码器的初始隐藏状态。这要求编码器和解码器的循环神经网络使用相同的隐藏层个数和隐藏单元个数。

在解码器的前向计算中，我们先通过刚刚介绍的注意力机制计算得到当前时间步的背景向量。由于解码器的输入来自输出语言的词索引，我们将输入通过词嵌入层得到表征，然后和背景向量在特征维连结。我们将连结后的结果与上一时间步的隐藏状态通过门控循环单元计算出当前时间步的输出与隐藏状态。最后，我们将输出通过全连接层变换为有关各个输出词的预测，形状为(批量大小, 输出词典大小)。

class Decoder(nn.Module):
    def __init__(self, vocab_size, embed_size, num_hiddens, num_layers,
                 attention_size, drop_prob=0):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(vocab_size, embed_size)
        self.attention = attention_model(2*num_hiddens, attention_size)
        # GRU的输入包含attention输出的c和实际输入, 所以尺寸是 num_hiddens+embed_size
        self.rnn = nn.GRU(num_hiddens + embed_size, num_hiddens, 
                          num_layers, dropout=drop_prob)
        self.out = nn.Linear(num_hiddens, vocab_size)

    def forward(self, cur_input, state, enc_states):
        """
        cur_input shape: (batch, )
        state shape: (num_layers, batch, num_hiddens)
        """
        # 使用注意力机制计算背景向量
        c = attention_forward(self.attention, enc_states, state[-1])
        # 将嵌入后的输入和背景向量在特征维连结, (批量大小, num_hiddens+embed_size)
        input_and_c = torch.cat((self.embedding(cur_input), c), dim=1) 
        # 为输入和背景向量的连结增加时间步维，时间步个数为1
        output, state = self.rnn(input_and_c.unsqueeze(0), state)
        # 移除时间步维，输出形状为(批量大小, 输出词典大小)
        output = self.out(output).squeeze(dim=0)
        return output, state

    def begin_state(self, enc_state):
        # 直接将编码器最终时间步的隐藏状态作为解码器的初始隐藏状态
        return enc_state

3.3 训练模型

我们先实现batch_loss函数计算一个小批量的损失。解码器在最初时间步的输入是特殊字符BOS。之后，解码器在某时间步的输入为样本输出序列在上一时间步的词，即强制教学。此外，我们在这里也使用掩码变量避免填充项对损失函数计算的影响。

def batch_loss(encoder, decoder, X, Y, loss):
    batch_size = X.shape[0]
    enc_state = encoder.begin_state()
    enc_outputs, enc_state = encoder(X, enc_state)
    # 初始化解码器的隐藏状态
    dec_state = decoder.begin_state(enc_state)
    # 解码器在最初时间步的输入是BOS
    dec_input = torch.tensor([out_vocab.stoi[BOS]] * batch_size)
    # 我们将使用掩码变量mask来忽略掉标签为填充项PAD的损失, 初始全1
    mask, num_not_pad_tokens = torch.ones(batch_size,), 0
    l = torch.tensor([0.0])
    for y in Y.permute(1,0): # Y shape: (batch, seq_len)
        dec_output, dec_state = decoder(dec_input, dec_state, enc_outputs)
        l = l + (mask * loss(dec_output, y)).sum()
        dec_input = y  # 使用强制教学
        num_not_pad_tokens += mask.sum().item()
        # EOS后面全是PAD. 下面一行保证一旦遇到EOS接下来的循环中mask就一直是0
        mask = mask * (y != out_vocab.stoi[EOS]).float()
    return l / num_not_pad_tokens

在训练函数中，我们需要同时迭代编码器和解码器的模型参数。

def train(encoder, decoder, dataset, lr, batch_size, num_epochs):
    """
    训练函数，用于训练编码器-解码器模型。

    参数:
    encoder (nn.Module): 编码器模型。
    decoder (nn.Module): 解码器模型。
    dataset (torch.utils.data.Dataset): 数据集对象。
    lr (float): 学习率。
    batch_size (int): 批量大小。
    num_epochs (int): 训练轮数。
    """
    # 定义编码器和解码器的优化器，使用Adam优化器。
    enc_optimizer = torch.optim.Adam(encoder.parameters(), lr=lr)
    dec_optimizer = torch.optim.Adam(decoder.parameters(), lr=lr)

    # 定义损失函数，交叉熵损失，但不进行求和，保留每个样本的损失。
    loss = nn.CrossEntropyLoss(reduction='none')

    # 创建数据迭代器，用于批量加载数据。
    data_iter = Data.DataLoader(dataset, batch_size, shuffle=True)

    # 开始训练循环
    for epoch in range(num_epochs):
        l_sum = 0.0  # 用于累加每个epoch的损失总和
        for X, Y in data_iter:
            enc_optimizer.zero_grad()  # 梯度清零
            dec_optimizer.zero_grad()  # 梯度清零
            l = batch_loss(encoder, decoder, X, Y, loss)  # 计算批量损失
            l.backward()  # 反向传播计算梯度
            enc_optimizer.step()  # 更新编码器参数
            dec_optimizer.step()  # 更新解码器参数
            l_sum += l.item()  # 累加损失值
        if (epoch + 1) % 10 == 0:
            print("epoch %d, loss %.3f" % (epoch + 1, l_sum / len(data_iter)))  # 每10个epoch打印损失平均值

接下来，创建模型实例并设置超参数。然后，我们就可以训练模型了。

# 定义编码器和解码器的隐藏单元数量及层数
embed_size, num_hiddens, num_layers = 64, 64, 2

# 定义注意力机制的大小、dropout概率、学习率、批大小和训练轮数
attention_size, drop_prob, lr, batch_size, num_epochs = 10, 0.5, 0.01, 2, 50

# 创建编码器对象，指定输入词汇表大小、嵌入维度、隐藏单元数量、层数和dropout概率
encoder = Encoder(len(in_vocab), embed_size, num_hiddens, num_layers,
                  drop_prob)

# 创建解码器对象，指定输出词汇表大小、嵌入维度、隐藏单元数量、层数、注意力大小和dropout概率
decoder = Decoder(len(out_vocab), embed_size, num_hiddens, num_layers,
                  attention_size, drop_prob)

# 调用训练函数，传入编码器、解码器、数据集、学习率、批大小和训练轮数
train(encoder, decoder, dataset, lr, batch_size, num_epochs)

3.4 预测不定长的序列

这里我们实现贪婪搜索。

def translate(encoder, decoder, input_seq, max_seq_len):
    # 将输入序列拆分为单词，并添加EOS和PAD，使其长度达到max_seq_len
    in_tokens = input_seq.split(' ')
    in_tokens += [EOS] + [PAD] * (max_seq_len - len(in_tokens) - 1)
    # 将单词转换为对应的索引，并构建成张量作为编码器的输入（batch=1）
    enc_input = torch.tensor([[in_vocab.stoi[tk] for tk in in_tokens]])
    # 获取编码器的初始状态
    enc_state = encoder.begin_state()
    # 对输入序列进行编码
    enc_output, enc_state = encoder(enc_input, enc_state)
    # 初始化解码器的输入为BOS
    dec_input = torch.tensor([out_vocab.stoi[BOS]])
    # 获取解码器的初始状态，传入编码器的最终状态
    dec_state = decoder.begin_state(enc_state)
    # 初始化输出序列
    output_tokens = []
    # 循环生成输出序列，最多为max_seq_len长度
    for _ in range(max_seq_len):
        # 解码器生成输出及更新状态
        dec_output, dec_state = decoder(dec_input, dec_state, enc_output)
        # 获取预测结果中概率最大的单词索引，并转换为对应的单词
        pred = dec_output.argmax(dim=1)
        pred_token = out_vocab.itos[int(pred.item())]
        # 如果预测结果为EOS，则停止生成序列
        if pred_token == EOS:
            break
        else:
            # 将预测结果添加到输出序列中，并更新解码器的输入为当前预测结果
            output_tokens.append(pred_token)
            dec_input = pred
    return output_tokens

简单测试一下模型。输入法语句子“ils regardent.”，翻译后的英语句子应该是“they are watching.”。

input_seq = 'ils regardent .'
translate(encoder, decoder, input_seq, max_seq_len)

3.5 评价翻译结果

评价机器翻译结果通常使用BLEU（Bilingual Evaluation Understudy）。对于模型预测序列中任意的子序列，BLEU考察这个子序列是否出现在标签序列中。

下面来实现BLEU的计算。

def bleu(pred_tokens, label_tokens, k):
    # 计算预测序列和参考序列的长度
    len_pred, len_label = len(pred_tokens), len(label_tokens)
    # 初始化BLEU分数
    score = math.exp(min(0, 1 - len_label / len_pred))
    
    # 计算1到k-gram的精确匹配
    for n in range(1, k + 1):
        num_matches, label_subs = 0, collections.defaultdict(int)
        
        # 统计参考序列中n-gram的出现次数
        for i in range(len_label - n + 1):
            label_subs[''.join(label_tokens[i: i + n])] += 1
        
        # 计算预测序列中与参考序列n-gram匹配的次数
        for i in range(len_pred - n + 1):
            if label_subs[''.join(pred_tokens[i: i + n])] > 0:
                num_matches += 1
                label_subs[''.join(pred_tokens[i: i + n])] -= 1
        
        # 计算BLEU分数的累积乘积
        score *= math.pow(num_matches / (len_pred - n + 1), math.pow(0.5, n))
    
    return score

接下来，定义一个辅助打印函数。

def score(input_seq, label_seq, k):
    # 使用encoder和decoder进行翻译，生成预测的token序列
    pred_tokens = translate(encoder, decoder, input_seq, max_seq_len)
    
    # 将标签序列按空格分割成token列表
    label_tokens = label_seq.split(' ')
    
    # 计算BLEU分数，并打印出结果
    print('bleu %.3f, predict: %s' % (bleu(pred_tokens, label_tokens, k),
                                      ' '.join(pred_tokens)))

预测正确则分数为1。

score('ils regardent .', 'they are watching .', k=2)

4. Transformer构建实例

本小节我们将构建一个日译中的Transformer代码实例。

4.1 导入库并下载数据集

我们将使用从 JParaCrawl 下载的日英并行数据集。

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

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.pop(5972)
# trainja.pop(5972)

4.2 准备分词器

与英语或其他字母语言不同，日语句子不包含空格来分隔单词。因此我们需要事先进行分词,这里的分词器是JParaCrawl提供的。

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

4.3 构建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)

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)

4.4 创建DataLoader object

这里还包括Batchsize的设置，调整它的大小可以改变模型复杂度。

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

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

以下代码内容包括模型参数，可以修改模型参数以配置模型，以3060Laptop版本（6G显存）为例，由于可用计算内存过小，只能将模型调到非常简单的程度。

SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(en_vocab)
EMB_SIZE = 256
NHEAD = 4
FFN_HID_DIM = 256
BATCH_SIZE = 4
NUM_ENCODER_LAYERS = 2
NUM_DECODER_LAYERS = 2
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)

4.6 开始训练

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

4.6 模型使用和保存

使用模型的代码：

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)

保存模型:

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


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