上一篇博客中，利用编码器-解码器结构，实现了机器翻译的功能。但是，其中，在注意力机制模块，只是一笔带过，因为在笔者看来，transformer 相比于基础的编码器-解码器结构，区别就是多头注意力机制以及自注意力机制的加入。况且，注意力机制对于编码器-解码器结构实际上的非必需的，上一篇只是为了训练结果能够有所提升才加入的。因此，笔者计划将对于注意力机制的内容留在本次进行介绍。

问题方面，机器翻译大家多少有所了解；文本的处理、编码器-解码器方面，笔者在之前的博客中也有介绍过。因此，本篇主要以介绍前面没有介绍过的内容为主。

一、模型介绍

词汇表、词嵌入、编码器-解码器结构在笔者之前的博客中已经介绍过了，因此，本次直接进入正题，讲讲注意力机制相关的内容。

1.1 注意力机制

1.1.1 什么是注意力机制

人眼是通过光线观察事物的，因此，理论上只要是光线进入你的瞳孔，都是你能看见的东西。但是当问你，你的面前有什么的时候，你往往只会回答你所注意的东西。而这个选择的过程，就是注意力机制。

例如，我们来做个实验：

当你看到这句话的时候

你的目光就会停在这句话上

你能够把这句话看的一清二楚

但是此时你是绝对看不清这句话的

这，就是，注意力机制！

对于机器学习而言，注意力机制的作用就是将原本普普通通的样本加上权重，让计算机能够接收到更多“被关注”的部分，而接收更少的“被忽略”的部分，从而实现更加精确的翻译结果。

具体的实现方式就是利用一种查询（Query）、键（Key）和值（Value）的方式实现的。

具体的方式为：在接收输入之后，将输出转化为键（非自主性提示）、查询（自主性提示）和值（感官输入）最终得到输出。这样得到的输入能够同时兼顾数据本身的特征和主动选择的结果，因此在一些任务中加入注意力机制，能够获得更优质的性能。

还是不好理解？没关系，我们举个例子：

你面前摆着 5 样东西：

白色的盘子，白色的刀叉，白色的勺子，白色的桌布，红色的咖啡杯

那么这五样东西就构成了值（感官输入）

由于在白色的东西中红色给你更大的冲击感，因此你会更容易注意咖啡杯

这就是键（非自主性提示）

但是你现在希望能够用勺子去搅拌咖啡，因此你会更容易注意到勺子

这就是查询（自主性提示）

1.1.2 如何使用注意力机制？

既然原理复杂，那就长话短说：键和查询是决定最终的注意力在什么位置，注意力的位置对应的值决定最终输出是什么。那我们可以制定一些比较简单的规则：例如，键和查询同时注意的部分给予更高的注意程度，注意程度更高的值在最终结果中权重更高。那这样，就可以设计一套规则：键和值两两相乘，那之后将值按照相乘的结果分类权重，例如使用 softmax 让值加权之后不会膨胀。由于样本的特性，键和值的长度可能不为 1，但是长度一定相等。此时利用点乘代替乘法即可。值的话也是同理。下面是这一过程的示意图：

上图所示就是注意力机制的实现方式示意图。其中，k-t表示的是键-值对，q表示的是查询，黄色方框表示的是归一化步骤，如 softmax 一类，o 表示的是输出。

上面这一过程就是从“我们希望注意力机制具有的功能”到“注意力机制的实现”的全过程。虽然最后的公式以及后续的优化看起来非常复杂，但是思路还是比较明确的，推导的过程也比较容易看懂。

观察最终的模型，以点乘为例，“查询”和“键”越接近，二者的点乘结果越大，最终“值”的权重也就越大。也就是输出结果越靠近“注意力”所注意的“值”。结果说明，我们构建的这个模型能够完成注意力机制的功能。这可能和一开始的表述不太一样，下面通过一个例子演示：

假设 t1=1, t2=0, t3=-1, k1=[1,0,0], k2=[0,1,0], k3=[0,0,1], 那这个时候如果查询 q1 主要关注点在 k1 附近，例如 q1 = [0.8,0.2,0], 那么就可以求解发现 k1*q1 = 0.8, k2*q1 = 0.2, k3*q1 = 0，为了便于演示先不加入 softmax，此时最终的结果就是 t1*0.8 + t2*0.2 + t3*0 = 0.8, 接近 t1。这样或许能够略微帮助理解一些。

1.1.3 多头注意力机制

上面的例子使用的是一个认为设定的注意力机制。但是实际的情况的不确定的，机器翻译的机理也是一个半黑盒的状态。那么如何让这一情况有所改变呢？为它加上代表映射关系的可以学习的全连接层。

不过既然都加上了映射关系，自然是不满足于一对一的映射，因为这样只会经过一次注意力汇聚的过程。所以，可以采用一对多的映射，从而使其经过多个不同的注意力汇聚过程。最后，为了维持输出的维数恒定，使用拼接+全连接层降维的方式对多个注意力汇聚的结果进行处理。这样的过程也就是多头注意力机制。

了解多头注意力机制的由来，了解结构的构成就可以了，更细节的内容，讲起来复杂，且了解到这一层已经不影响读者理解模型的功能了。如果想要进一步了解原理，可以选择去看一看 Transformer 的原文。

1.1.4 自注意力机制

就像是 Transformer 的论文标题一样，Attention is all you need。Transformer 结构不仅在解码器中运用了注意力机制，同样在循环层中使用了它。不同于原本的，按照时间步依次读取序列元素的方式，Transformer 使用了一种注意力机制去并行的读取并记录句子的特征。这种注意力机制中，键、值和查询三者的来源都是同一个句子，因此也叫做自注意力机制。

对于实现而言，自注意力机制可以使用 1.1.3 节中的多头注意力机制实现。只需要将三者同时输入同一个样本序列就能够实现。不过需要注意的是，此时的键、值、查询并不是 1.1.3 中所示的内容，而是 1.1.3 图中的键、值、查询分别加上一层线性全连接网络。这种做法在数学上就像是原本的输入乘上对应的矩阵得到了键、值和查询，因此你也许会在看到一些数学公式的时候感到害怕。但是不必担心，大部分情况下，数学公式中的矩阵对于机器学习的过程而言，都可以用线性全连接网络进行实现。

1.2 位置编码

1.2.1 什么是位置编码

位置编码是一个能够通过目标元素的横纵坐标生成的，仅与元素在坐标系中的位置有关且具有唯一性的，能够表示任意两个元素之间距离的一种给元素编码表示位置的方式。加入位置编码之后，就相当于元素除了原本的数值属性，还增加了位置属性。这一点在一些需要考虑元素位置或者元素的排列顺序的任务中会有比较不错的效果。

$PE_{pos,2i}=\sin(\frac{pos}{10000^{\frac{2i}{d}}}) \\ PE_{pos,2i+1}=\cos(\frac{pos}{10000^{\frac{2i}{d}}})$

例如，对于机器翻译任务，上述公式中，pos 代表词在句子中第几个位置，2i 或者 2i+1 表示的是词嵌入之后，词向量的第几个维度，d 表示的则是词嵌入之后，词向量的维度大小。这样能够对词嵌入之后的每个数值都赋予一个位置。没听错，在 Transformer 中，位置编码的编码对象不是句子也不是词，而是词向量某个维度的数值。这听起来匪夷所思，后面会进一步解释。

1.2.2 为什么使用位置编码

在 1.1.4 节中有提到，Transformer 模型使用自注意力机制实现对于句子特征的并行传入。并行传入是自然有着快速、方便一类的优势，但是也是由于并行传入，系统丢失了最重要的顺序信息。这样的话，“我是你爸爸”和“你是我爸爸”两句截然不同的话，在模型看来可能是同样的一句话，这并不是我们想要的。因此，Transformer 模型在词嵌入之后，自注意力机制之前，加上了一层位置编码，使得即便是并行输入的数据，计算机也能够读取每个词乃至词向量每个维度的位置信息。这对于从逻辑层面提升模型的精度而言非常重要。

1.3 Transformer 模型

这个模型图我想各位读者应该在其他地方多少有见过一面：

ps:此图并非原创，只做讲解使用。

大题的模型和原本的编码器-解码器结构相类似，但是多了一些简单编码器-解码器结构不具有的细节。

下面笔者会着重介绍这些改动，除了前面已经介绍的自注意力机制（使用多头注意力机制实现，Multi-Head Attention）和位置编码（Positional Encoding），以及笔者之前的文章中介绍过的词嵌入部分（Embedding）、逐位前馈神经网络（上一篇中介绍的 GRU 的定位，Feed Forward）等。

这里其实存在一个非常有意思的问题：

词嵌入之后，需要添加位置编码到输入的数据中。但是这里使用的是简单的按位相加。原作者在这一部分似乎只提到了“两者维度相同，因此可以相加”，而没有说为什么相加对于位置编码与词向量的组合而言是一个有效的方式。

这是一个笔者的思考，欢迎对这一部分有更多了解的大佬在评论区为我解惑，也欢迎同样感到疑惑的小伙伴私信交流心得。

1.3.1 残差连接（Residual Connection）

可以观察到，Transformer 在每一个层的侧面，都设计了一个上一层的输入可以直接去到下一层的“捷径”，而添加这个“捷径”的做法一般被称作残差链接。做法也很简单，在前向传播的过程中，将原本的输出y = model(x)改写成 y = (x + model(x)) 就可以了。

Q：为什么要使用残差连接？

A：为了保持反向传播过程的通畅。对于机器学习的参数更新而言，目前主流的方式仍然是使用反向传播的方式。但是这种方式带来的就是梯度爆炸和梯度消失的可能性。残差连接就是改善这一问题的一个方式。如果对这一块的原理不太明了的，可以去看一看 ResNet 模型的讲解，这里只是提到，并不希望以之为主。

1.3.2 加法和归一化（Add & Norm）

加法其实是和残差连接相关的操作，就是将两条路线的结果相加，因此这里主要介绍归一化。

这里的归一化一般指的是层归一化。这是由于加入了残差连接，输入和输出之间的关系不对等。如果不加入层归一化，容易在层层累加之下出现诸如训练结果不收敛一类的问题。因此，这里需要进行归一化。

1.3.3 带有掩码的多头注意力层（Masked Multi-Head Attention）

与编码器不同，解码器是通过将上一时间步解码器的输出作为下一个解码器的输入来进行结果的预测的，因此，在训练的过程中，只能输入对应时间步之前的变量，即，加入掩码机制，先遮住后面部分，只考虑前面部分。

1.3.4 编码器-解码器注意力层（Multi-Head Attention）

解码器和编码器还有一点不同，就是在两个子层之间插入了第三个子层，用于接收编码器的输出。这一部分的功能实际上就相当于简单编码器-解码器结构中接收编码器的输出作为隐状态这一步骤，但是使用多头注意力机制作为实现方式。对于这一层而言，查询为解码器前面步骤的输出，键和值来自编码器的输出。这样就能将编码器的结果挂载到解码器上。

Q：为什么编码器连接键和值，解码器前面的步骤连接查询？

A：这一过程可以理解为我们查字典的过程。查询表示的是要查的内容，而键-值对是记录下来的对照规则。也就是说，解码器的工作是，按照编码器得到的规则，依次查询得到结果输出的过程，因此，“查询”是解码器该干的，而“键-值”对的内容，应当是编码器得到的结果。

当然，这种想法只能用于理解，实际的原理可能会更复杂。

模型部分至此介绍完毕，下面准备进行代码的编写。

二、数据集介绍

本次实验使用的数据来源如下：

http://www.kecl.ntt.co.jp/icl/lirg/jparacrawlhttp://www.kecl.ntt.co.jp/icl/lirg/jparacrawl本次使用的是其中的英语-日语数据集，但是为了便于操作，这里选择将英语数据翻译为中文数据，从而方便结果的展示与实验者的理解。

三、实验环境

笔者使用的是MacBook16Pro 搭载 MacOS 系统。由于某些原因，不能装载 N 卡和 CUDA，因此使用 SSH 的方式挂载云端算力进行。以下是一些云端的环境等信息：

GPU：

CUDA 11.3

编译软件：

Pycharm

python 包版本(python=3.8)：

d2l 0.17.5
ipython 8.4.0
matplotlib 3.5.1
numpy 1.21.5
pandas 1.2.4
Pillow 9.1.1
sentencepiece 0.1.92
torch 1.11.0+cu113
torchinfo 1.8.0
torchtext 0.6.0
torchvision 0.12.0+cu113
tqdm 4.61.2

如果有其他不确定的部分欢迎私信交流。

四、代码部分

原理介绍完毕，软件版本问题解决完毕，所以下面开始代码的编写工作。

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

包的版本在上面没有列出来的，一般而言，按照 python 版本让 pip 自动选择合适的版本即可。

4.2 数据读取与数据集建立

这一部分和前面的基本一致，都是读取、分词、生成词汇表、返回「索引」化之后的数据集等几个步骤。原理部分应当都已经介绍过，这里就不展开了。基本的思路和编码器解码器一致，可以利用上一次的思路进行。不过这一次使用了预训练的分词器，因此在生成数据集的时候会更简单一些。

def build_vocab(sentences, tokenizer):
    """
    生成词汇表

    :param sentences: 句子
    :param tokenizer: 对应语言的分词器
    :return:
    """
    counter = Counter()
    for sentence in sentences:
        counter.update(tokenizer.encode(sentence, out_type=str))
    return Vocab(counter, specials=['<unk>', '<pad>', '<bos>', '<eos>'])

# 读取数据集
df = pd.read_csv('../data14/zh-ja.bicleaner05.txt', sep='\\t', engine='python', header=None)

# 初步划分为两个语言的原始数据集
train_ch = df[2].values.tolist()#[:10000]
train_ja = df[3].values.tolist()#[:10000]

# 读取预训练的分词器
ch_tokenizer = spm.SentencePieceProcessor(model_file='../data14/spm.en.nopretok.model')
ja_tokenizer = spm.SentencePieceProcessor(model_file='../data14/spm.ja.nopretok.model')

# 构建词汇表
ja_vocab = build_vocab(train_ja, ja_tokenizer)

# 取出特殊值的索引备用
PAD_IDX = ja_vocab['<pad>']
BOS_IDX = ja_vocab['<bos>']
EOS_IDX = ja_vocab['<eos>']


def data_process(ja, en, ja_vocab, ja_tokenizer, en_vocab, en_tokenizer):
    """
    数据处理
        将数据集按照词汇表进行「索引」化，形成处理之后的数据集

    :param ja: 日文原始数据集
    :param en: 中文原始数据集
    :param ja_vocab: 日文词汇表
    :param ja_tokenizer: 日文分词器
    :param en_vocab: 中文词汇表
    :param en_tokenizer: 中文分词器
    :return:
    """
    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


def generate_batch(data_batch):
    """
    定义 DataLoader 批量读取数据的方法
    
    :param data_batch: 数据批次（一个 DataLoader 中的中间变量）
    :return: 
    """
    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

# 批次大小
BATCH_SIZE = 8
# 生成数据加载器
train_iter = DataLoader(train_data, batch_size=BATCH_SIZE,
                        shuffle=True, collate_fn=generate_batch)

到数据加载器这一步，基本就能够当作训练输入的格式了。至此，我们完成了数据的读取和数据集的建立。

4.3 Transformer 模型构建

按照Transformer 的结构示意图进行构建即可。大部分都是已经成型的包装好的模型，直接使用即可。词嵌入以及位置编码部分，按照原理转化为代码实现即可制作，也非常简单。

class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int,
                 emb_size: int, num_heads: int, src_vocab_size: int, tgt_vocab_size: int,
                 dim_feedforward:int = 512, dropout:float = 0.1):
        """
        定义部分，主要介绍包含了哪些结构

        :param num_encoder_layers: 编码器层数
        :param num_decoder_layers: 解码器层数
        :param emb_size: 词嵌入映射词向量的空间维数
        :param num_heads: 自注意力机制中的头数
        :param src_vocab_size: 原语言词汇表
        :param tgt_vocab_size: 目的语言词汇表
        :param dim_feedforward: 前馈神经网络中隐藏层的维度
        :param dropout: dropout 比例
        """
        super(Seq2SeqTransformer, self).__init__()
        # 定义编码器层
        encoder_layer = TransformerEncoderLayer(d_model=emb_size, nhead=num_heads,
                                                dim_feedforward=dim_feedforward)
        # 编码器层叠加为编码器
        self.transformer_encoder = TransformerEncoder(encoder_layer, num_layers=num_encoder_layers)
        # 定义解码器层
        decoder_layer = TransformerDecoderLayer(d_model=emb_size, nhead=num_heads,
                                                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):
        """
        前向传播过程

        :param src: 编码器的输入序列
        :param trg: 解码器的输入序列
        :param src_mask: 编码器输入序列的掩码
        :param tgt_mask: 解码器输出序列的掩码
        :param src_padding_mask: 编码器输入的扩充键掩码
        :param tgt_padding_mask: 解码器输入的扩充键掩码
        :param memory_key_padding_mask: 编码器的最后一层输出序列的扩充键掩码
        :return: 前向传播的结果
        """
        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):
        """
        单独的编码器部分

        :param src: 编码器的输入序列
        :param src_mask: 源序列的掩码
        :return: 编码器最后一层的输出
        """
        return self.transformer_encoder(self.positional_encoding(
                            self.src_tok_emb(src)), src_mask)

    def decode(self, tgt: Tensor, memory: Tensor, tgt_mask: Tensor):
        """
        单独的解码器部分

        :param tgt: 解码器的输入序列
        :param memory: 编码器的最后一层输出
        :param tgt_mask: 目标序列的掩码
        :return: 解码器的输出结果
        """
        return self.transformer_decoder(self.positional_encoding(
                          self.tgt_tok_emb(tgt)), memory,
                          tgt_mask)


class PositionalEncoding(nn.Module):
    """
    位置编码
        - 用于将 每个词都经过词嵌入 的向量语句，每一个数值所在的位置进行编码
        - 例如：如果词嵌入之后向量是 5 维的，也就是 d=5，第一个单词的第一个维度对应的索引就是 p=0, i=0
        - 对应的位置编码就是：pe = sin(p / 10000 ** (i/d)）
        -
    """
    def __init__(self, emb_size: int, dropout, maxlen: int = 5000):
        """
        属性定义部分

        :param emb_size: 设定的向量空间的维数
        :param dropout: 随机失活的比例
        :param maxlen: 限制最大的序列长度
        """
        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))
        # 步长为 2 的切片，从而区分奇偶
        pos_embedding[:, 0::2] = torch.sin(pos * den)
        pos_embedding[:, 1::2] = torch.cos(pos * den)
        # 为了与后面 token_embedding 的维度保持一致，释放了一个维度
        pos_embedding = pos_embedding.unsqueeze(-2)

        self.dropout = nn.Dropout(dropout)
        self.register_buffer('pos_embedding', pos_embedding)

    def forward(self, token_embedding: Tensor):
        """
        前向传播部分

        :param token_embedding: 词嵌入之后的单词
        :return:
        """
        return self.dropout(token_embedding +
                            self.pos_embedding[:token_embedding.size(0),:])


class TokenEmbedding(nn.Module):
    """
    词嵌入 类
        将已经「索引」化之后的语句进行向量化
    """
    def __init__(self, vocab_size: int, emb_size):
        """
        模型定义

        :param vocab_size: 词汇表的大小
        :param 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)

需要的 seq2seqTransformer 类定义好了，然后就是将其实例化：

# 参数的含义在模型中都介绍了，这里就不再讲一遍了
SRC_VOCAB_SIZE = len(ja_vocab)
TGT_VOCAB_SIZE = len(ch_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, NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE,
                                 FFN_HID_DIM)

4.4 训练以及验证过程

就是经典的训练循环体，没有什么比较独特的点需要介绍。

def generate_square_subsequent_mask(sz):
    """
    根据目标序列长度创建对应的掩码

    :param sz: 目标序列长度
    :return:
    """
    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):
    """
    创建编码器输入和解码器输入序列的掩码

    :param src: 编码器输入
    :param tgt: 解码器输入
    :return:
    """
    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


def train_epoch(model, train_iter, loss_fn, optimizer, epoch_idx, num_epoch):
    """
    训练过程

    :param model: Transformer 模型
    :param train_iter: 训练集
    :param loss_fn: 损失函数
    :param optimizer: 优化器
    :param epoch_idx: 训练轮次
    :param num_epoch: 训练总轮次
    :return: 训练集上的损失大小
    """
    model.train()
    losses = 0
    for idx, (src, tgt) in tqdm(enumerate(train_iter),
                                desc=f"epoch {epoch_idx}/{num_epoch}",
                                total=len(train_iter)):
        src = src.to(device)
        tgt = tgt.to(device)

        tgt_input = tgt[:-1, :]

        # 创建掩码
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = hp.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, loss_fn):
    """
    验证模型

    :param model: Transformer 模型
    :param val_iter: 验证集
    :param loss_fn: 损失函数
    :return: 验证集上的损失大小
    """

    # 验证模式
    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 = hp.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 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
)

for epoch in range(1, NUM_EPOCHS+1):
    start_time = time.time()
    train_loss = mt.train_epoch(transformer, train_iter, loss_fn, optimizer, epoch, NUM_EPOCHS )
    end_time = time.time()
    print((f"Epoch: {epoch}, Train loss: {train_loss:.3f}, "
           f"Epoch time = {(end_time - start_time):.3f}s"))

训练过程如下（真是挺长一段时间）：

epoch 1/16: 100%|██████████| 10487/10487 [05:40<00:00, 30.79it/s]
Epoch: 1, Train loss: 4.475, Epoch time = 340.585s
epoch 2/16: 100%|██████████| 10487/10487 [05:38<00:00, 30.94it/s]
Epoch: 2, Train loss: 3.487, Epoch time = 338.939s
epoch 3/16: 100%|██████████| 10487/10487 [05:38<00:00, 31.00it/s]
Epoch: 3, Train loss: 3.080, Epoch time = 338.302s
epoch 4/16: 100%|██████████| 10487/10487 [05:41<00:00, 30.70it/s]
Epoch: 4, Train loss: 2.785, Epoch time = 341.627s
epoch 5/16: 100%|██████████| 10487/10487 [05:36<00:00, 31.14it/s]
Epoch: 5, Train loss: 2.567, Epoch time = 336.797s
epoch 6/16: 100%|██████████| 10487/10487 [05:35<00:00, 31.27it/s]
Epoch: 6, Train loss: 2.400, Epoch time = 335.404s
epoch 7/16: 100%|██████████| 10487/10487 [05:36<00:00, 31.21it/s]
Epoch: 7, Train loss: 2.292, Epoch time = 336.059s
epoch 8/16: 100%|██████████| 10487/10487 [05:37<00:00, 31.06it/s]
Epoch: 8, Train loss: 2.200, Epoch time = 337.592s
epoch 9/16: 100%|██████████| 10487/10487 [05:35<00:00, 31.23it/s]
Epoch: 9, Train loss: 2.117, Epoch time = 335.854s
epoch 10/16: 100%|██████████| 10487/10487 [05:37<00:00, 31.07it/s]
Epoch: 10, Train loss: 2.046, Epoch time = 337.523s
epoch 11/16: 100%|██████████| 10487/10487 [05:37<00:00, 31.06it/s]
Epoch: 11, Train loss: 1.985, Epoch time = 337.611s
epoch 12/16: 100%|██████████| 10487/10487 [05:37<00:00, 31.11it/s]
Epoch: 12, Train loss: 1.933, Epoch time = 337.137s
epoch 13/16: 100%|██████████| 10487/10487 [05:38<00:00, 31.00it/s]
Epoch: 13, Train loss: 1.887, Epoch time = 338.309s
epoch 14/16: 100%|██████████| 10487/10487 [05:36<00:00, 31.18it/s]
Epoch: 14, Train loss: 1.846, Epoch time = 336.289s
epoch 15/16: 100%|██████████| 10487/10487 [05:36<00:00, 31.12it/s]
Epoch: 15, Train loss: 1.808, Epoch time = 336.935s
epoch 16/16: 100%|██████████| 10487/10487 [05:37<00:00, 31.04it/s]
Epoch: 16, Train loss: 1.773, Epoch time = 337.821s

4.5 结果检验与模型保存

检验就是观察使用模型实际的翻译效果如何，因此，这里定义一个翻译功能的函数：

def greedy_decode(model, src, src_mask, max_len, start_symbol):
    """
    贪婪算法解码器
        用于翻译过程，即不考虑全局最优，只考虑局部最优，一步步得到最终结果

    :param model: Transformer 模型
    :param src: 「索引」化之后的日文序列
    :param src_mask: 日文序列对应掩码
    :param max_len: 最大句长
    :param start_symbol: 开始标志
    :return:
    """
    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):
    """
    利用训练后的结果进行翻译的过程

    :param model: Transformer 模型
    :param src: 日文序列
    :param src_vocab: 日文对应词汇表
    :param tgt_vocab: 中文对应词汇表
    :param src_tokenizer: 日文分词器
    :return:
    """
    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>", "")

试试看翻译的结果：

print(translate(transformer, "HSコード 8515 はんだ付け用、ろう付け用又は溶接用の機器(電気式(電気加熱ガス式を含む。)", ja_vocab, ch_vocab, ja_tokenizer))

▁H S 还 是 焊 接 机 , 包 括 电 气 加 热 和 焊 接 设 备 ( 包 括 电 气 加 热 ) 。

可以看到，经过了长时间的训练过程，翻译的结果算是比较接近了。

如果认可此时的训练结果，可以将此时的模型保存下来：

file = open('vocab/ch_vocab.pkl', 'wb')
# dump information to that file
pickle.dump(ch_vocab, file)
file.close()

file = open('vocab/ja_vocab.pkl', 'wb')
pickle.dump(ja_vocab, file)
file.close()

torch.save(transformer.state_dict(), 'weight/inference_model')

代码到这里就结束了。如果有哪些部分有问题的话，欢迎给我留言，笔者在这里谢谢各位了。