Pytorch入门实战（5）：基于nn.Transformer实现机器翻译（英译汉）

使用Google Colab运行(open In Colab)

源码地址

本文涉及知识点

1. nn.Transformer的使用
2. Transformer源码解读 (了解即可)
3. Pytorch中DataLoader和Dataset的基本用法
4. subword基本概念
5. Masked-Attention的机制和原理
6. Pytorch自定义损失函数
7. Pytorch使用TensorBoard

本文内容

本文将使用Pytorch提供的nn.Transformer实现英文到中文的机器翻译任务。对nn.Transformer的讲解，可以参考我的另一篇博文Pytorch中 nn.Transformer的使用详解与Transformer的黑盒讲解，建议先学习该文的CopyTask任务，然后再来看该篇就容易多了。

本篇内容要求对Transformer有一定的了解，尤其是Transformer的入参出参、训练方式、推理方式和Mask部分。这些内容在上面的本文涉及知识点中找到。

本篇源码可以在该github项目中找到。

本篇最终效果：

translate("Alright, this project is finished. Let's see how good this is")
1

'好吧，这个项目完成了。让我们看看这是多好的。'
1

这是我训练了10个小时的效果。（1个epoch都没跑完，loss其实还能降）

环境配置

本文主要使用到的环境如下：

torch>=1.11.0
tokenizers==0.12.1
torchtext==0.12.0
tensorboard==2.8.0
1
2
3
4

首先我们需要导入本文需要用到的包：

import os
import math

import torch
import torch.nn as nn
# hugging face的分词器，github地址：https://github.com/huggingface/tokenizers
from tokenizers import Tokenizer
# 用于构建词典
from torchtext.vocab import build_vocab_from_iterator
from torch.utils.data import Dataset
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
from torch.nn.functional import pad, log_softmax
from pathlib import Path
from tqdm import tqdm
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

下载数据集。数据集包括两个文件，train.en和train.zh。这两个都是文本文件，里面存放了英文和中文的句子。

本文使用的是AI Challenger Translation 2017数据集。这里我简单进行了整理，只使用了其中的train.en和train.zh文件（简单起见，本文就不使用验证集了），同时我也将初始化的缓存文件放在了其中，直接解压即可。

百度网盘链接：链接：https://pan.baidu.com/s/1i9Ykz3YVdmKzQ0oKecdvaQ?pwd=4usf 提取码：4usf

如果你不想使用我缓存好的文件，可以将*.pt文件删除，或设置use_cache=False

定义一些全局配置，例如工作目录，训练时的batch_size，epoch等。

# 工作目录，缓存文件盒模型会放在该目录下
work_dir = Path("./dataset")
# 训练好的模型会放在该目录下
model_dir = Path("./drive/MyDrive/model/transformer_checkpoints")
# 上次运行到的地方，如果是第一次运行，为None，如果中途暂停了，下次运行时，指定目前最新的模型即可。
model_checkpoint = None # 'model_10000.pt'

# 如果工作目录不存在，则创建一个
if not os.path.exists(work_dir):
    os.makedirs(work_dir)

# 如果工作目录不存在，则创建一个
if not os.path.exists(model_dir):
    os.makedirs(model_dir)

# 英文句子的文件路径
en_filepath = './dataset/train.en'
# 中文句子的文件路径
zh_filepath = './dataset/train.zh'


# 定义一个获取文件行数的方法。
def get_row_count(filepath):
    count = 0
    for _ in open(filepath, encoding='utf-8'):
        count += 1
    return count


# 英文句子数量
en_row_count = get_row_count(en_filepath)
# 中文句子数量
zh_row_count = get_row_count(zh_filepath)
assert en_row_count == zh_row_count, "英文和中文文件行数不一致！"
# 句子数量，主要用于后面显示进度。
row_count = en_row_count

# 定义句子最大长度，如果句子不够这个长度，则填充，若超出该长度，则裁剪
max_length = 72
print("句子数量为：", en_row_count)
print("句子最大长度为：", max_length)

# 定义英文和中文词典，都为Vocab类对象，后面会对其初始化
en_vocab = None
zh_vocab = None

# 定义batch_size，由于是训练文本，占用内存较小，可以适当大一些
batch_size = 64
# epochs数量，不用太大，因为句子数量较多
epochs = 10
# 多少步保存一次模型，防止程序崩溃导致模型丢失。
save_after_step = 5000

# 是否使用缓存，由于文件较大，初始化动作较慢，所以将初始化好的文件持久化
use_cache = True

# 定义训练设备
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

print("batch_size:", batch_size)
print("每{}步保存一次模型".format(save_after_step))
print("Device:", device)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62

句子数量为： 10000000
句子最大长度为： 72
batch_size: 64
每5000步保存一次模型
Device: cuda
1
2
3
4
5

数据预处理

本章进行数据处理，主要过程有：

1. 构造英文词典和中文词典，其中英文采用subword方式，中文直接按字进行分词。
2. 构造Dataset和Dataloader，在其中对文本进行文本转数字(index)和值填充。

文本分词与构造词典

本文针对英文分词使用了subword的方式（subword相关概念）。分词器使用的是hugging face的bert模型，该分词器使用简单，不需要刻意学习，直接看本文就能看懂。

接下来来构造英文词典：

# 加载基础的分词器模型，使用的是基础的bert模型。`uncased`意思是不区分大小写
tokenizer = Tokenizer.from_pretrained("bert-base-uncased")


def en_tokenizer(line):
    """
    定义英文分词器，后续也要使用
    :param line: 一句英文句子，例如"I'm learning Deep learning."
    :return: subword分词后的记过，例如：['i', "'", 'm', 'learning', 'deep', 'learning', '.']
    """
    # 使用bert进行分词，并获取tokens。add_special_tokens是指不要在结果中增加‘<bos>’和`<eos>`等特殊字符
    return tokenizer.encode(line, add_special_tokens=False).tokens
1
2
3
4
5
6
7
8
9
10
11
12

如果你使用的是高版本，且上面这段报错了， 可以修改为如下代码：

from transformers import BertTokenizer
tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
def en_tokenizer(line):
   return tokenizer.convert_ids_to_tokens(tokenizer.encode(line,  add_special_tokens=False))
1
2
3
4

我们来测试一下英文分词器：

print(en_tokenizer("I'm a English tokenizer."))
1

['i', "'", 'm', 'a', 'english', 'token', '##izer', '.']
1

上面的分词结果中，tokenizer被拆成了两个subword：token和##izer。其中##表示这个词前面需要词与其连接。

接下来开始正式开始构造词典，我们先定义一个yield函数，来产生一个可迭代的分词结果：

def yield_en_tokens():
    """
    每次yield一个分词后的英文句子，之所以yield方式是为了节省内存。
    如果先分好词再构造词典，那么将会有大量文本驻留内存，造成内存溢出。
    """
    file = open(en_filepath, encoding='utf-8')
    print("-------开始构建英文词典-----------")
    for line in tqdm(file, desc="构建英文词典", total=row_count):
        yield en_tokenizer(line)
    file.close()
1
2
3
4
5
6
7
8
9
10

# 指定英文词典缓存文件路径
en_vocab_file = work_dir / "vocab_en.pt"
# 如果使用缓存，且缓存文件存在，则加载缓存文件
if use_cache and os.path.exists(en_vocab_file):
    en_vocab = torch.load(en_vocab_file, map_location="cpu")
# 否则就从0开始构造词典
else:
    # 构造词典
    en_vocab = build_vocab_from_iterator(
        # 传入一个可迭代的token列表。例如[['i', 'am', ...], ['machine', 'learning', ...], ...]
        yield_en_tokens(),
        # 最小频率为2，即一个单词最少出现两次才会被收录到词典
        min_freq=2,
        # 在词典的最开始加上这些特殊token
        specials=["<s>", "</s>", "<pad>", "<unk>"],
    )
    # 设置词典的默认index，后面文本转index时，如果找不到，就会用该index填充
    en_vocab.set_default_index(en_vocab["<unk>"])
    # 保存缓存文件
    if use_cache:
        torch.save(en_vocab, en_vocab_file)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

# 打印一下看一下效果
print("英文词典大小:", len(en_vocab))
print(dict((i, en_vocab.lookup_token(i)) for i in range(10)))
1
2
3

英文词典大小: 27584
{0: '<s>', 1: '</s>', 2: '<pad>', 3: '<unk>', 4: '.', 5: ',', 6: 'the', 7: "'", 8: 'i', 9: 'you'}
1
2

---

接着我们来构建中文词典，中文词比较多，容易产生OOV问题。一个简单的方式就是不分词，直接将每个字作为一个token，这么做对于中文来说是合理的，因为中文将一个词拆成字大多也能具备其含义，例如：单词一词，即使拆成单和词也能有原本的意思（单个词）。

构造中文词典和英文的类似：

def zh_tokenizer(line):
    """
    定义中文分词器
    :param line: 中文句子，例如：机器学习
    :return: 分词结果，例如['机','器','学','习']
    """
    return list(line.strip().replace(" ", ""))


def yield_zh_tokens():
    file = open(zh_filepath, encoding='utf-8')
    for line in tqdm(file, desc="构建中文词典", total=row_count):
        yield zh_tokenizer(line)
    file.close()
1
2
3
4
5
6
7
8
9
10
11
12
13
14

zh_vocab_file = work_dir / "vocab_zh.pt"
if use_cache and os.path.exists(zh_vocab_file):
    zh_vocab = torch.load(zh_vocab_file, map_location="cpu")
else:
    zh_vocab = build_vocab_from_iterator(
        yield_zh_tokens(),
        min_freq=1,
        specials=["<s>", "</s>", "<pad>", "<unk>"],
    )
    zh_vocab.set_default_index(zh_vocab["<unk>"])
    torch.save(zh_vocab, zh_vocab_file)
1
2
3
4
5
6
7
8
9
10
11

# 打印看一下效果
print("中文词典大小:", len(zh_vocab))
print(dict((i, zh_vocab.lookup_token(i)) for i in range(10)))
1
2
3

中文词典大小: 8280
{0: '<s>', 1: '</s>', 2: '<pad>', 3: '<unk>', 4: '。', 5: '的', 6: '，', 7: '我', 8: '你', 9: '是'}
1
2

Dataset and Dataloader

构造词典就可以来定义Dataset了。Dataset每次返回一个句子对儿，例如: ([6, 8, 93, 12, ..], [62, 891, ...])，第一个是英文句子，第二个是中文句子。

class TranslationDataset(Dataset):

    def __init__(self):
        # 加载英文tokens
        self.en_tokens = self.load_tokens(en_filepath, en_tokenizer, en_vocab, "构建英文tokens", 'en')
        # 加载中文tokens
        self.zh_tokens = self.load_tokens(zh_filepath, zh_tokenizer, zh_vocab, "构建中文tokens", 'zh')

    def __getitem__(self, index):
        return self.en_tokens[index], self.zh_tokens[index]

    def __len__(self):
        return row_count

    def load_tokens(self, file, tokenizer, vocab, desc, lang):
        """
        加载tokens，即将文本句子们转换成index们。
        :param file: 文件路径，例如"./dataset/train.en"
        :param tokenizer: 分词器，例如en_tokenizer函数
        :param vocab: 词典, Vocab类对象。例如 en_vocab
        :param desc: 用于进度显示的描述，例如：构建英文tokens
        :param lang: 语言。用于构造缓存文件时进行区分。例如：’en‘
        :return: 返回构造好的tokens。例如：[[6, 8, 93, 12, ..], [62, 891, ...], ...]
        """

        # 定义缓存文件存储路径
        cache_file = work_dir / "tokens_list.{}.pt".format(lang)
        # 如果使用缓存，且缓存文件存在，则直接加载
        if use_cache and os.path.exists(cache_file):
            print(f"正在加载缓存文件{cache_file}, 请稍后...")
            return torch.load(cache_file, map_location="cpu")

        # 从0开始构建，定义tokens_list用于存储结果
        tokens_list = []
        # 打开文件
        with open(file, encoding='utf-8') as file:
            # 逐行读取
            for line in tqdm(file, desc=desc, total=row_count):
                # 进行分词
                tokens = tokenizer(line)
                # 将文本分词结果通过词典转成index
                tokens = vocab(tokens)
                # append到结果中
                tokens_list.append(tokens)
        # 保存缓存文件
        if use_cache:
            torch.save(tokens_list, cache_file)

        return tokens_list
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49

dataset = TranslationDataset()
1

正在加载缓存文件dataset/tokens_list.en.pt, 请稍后...
正在加载缓存文件dataset/tokens_list.zh.pt, 请稍后...
1
2

定义好dataset后，我们来简单的看一下：

print(dataset.__getitem__(0))
1

([11, 2730, 12, 554, 19, 17210, 18077, 27, 3078, 203, 57, 102, 18832, 3653], [12, 40, 1173, 1084, 3169, 164, 693, 397, 84, 100, 14, 5, 1218, 2397, 535, 67])
1

Dataset中并不包含<bos>和<eos>，这个动作和填充是在dataloader中完成的。

接下来开始定义Dataloader。

在定义Dataloader之前，我们需要先定义好collate_fn，因为我们dataset返回的字段并不能很好的组成batch，并且需要进一步处理，这些操作的都是在collate_fn中完成。

def collate_fn(batch):
    """
    将dataset的数据进一步处理，并组成一个batch。
    :param batch: 一个batch的数据，例如：
                  [([6, 8, 93, 12, ..], [62, 891, ...]),
                  ....
                  ...]
    :return: 填充后的且等长的数据，包括src, tgt, tgt_y, n_tokens
             其中src为原句子，即要被翻译的句子
             tgt为目标句子：翻译后的句子，但不包含最后一个token
             tgt_y为label：翻译后的句子，但不包含第一个token，即<bos>
             n_tokens：tgt_y中的token数，<pad>不计算在内。
    """

    # 定义'<bos>'的index，在词典中为0，所以这里也是0
    bs_id = torch.tensor([0])
    # 定义'<eos>'的index
    eos_id = torch.tensor([1])
    # 定义<pad>的index
    pad_id = 2

    # 用于存储处理后的src和tgt
    src_list, tgt_list = [], []

    # 循环遍历句子对儿
    for (_src, _tgt) in batch:
        """
        _src: 英语句子，例如：`I love you`对应的index
        _tgt: 中文句子，例如：`我 爱 你`对应的index
        """

        processed_src = torch.cat(
            # 将<bos>，句子index和<eos>拼到一块
            [
                bs_id,
                torch.tensor(
                    _src,
                    dtype=torch.int64,
                ),
                eos_id,
            ],
            0,
        )
        processed_tgt = torch.cat(
            [
                bs_id,
                torch.tensor(
                    _tgt,
                    dtype=torch.int64,
                ),
                eos_id,
            ],
            0,
        )

        """
        将长度不足的句子进行填充到max_padding的长度的，然后增添到list中

        pad：假设processed_src为[0, 1136, 2468, 1349, 1]
             第二个参数为: (0, 72-5)
             第三个参数为：2
        则pad的意思表示，给processed_src左边填充0个2，右边填充67个2。
        最终结果为：[0, 1136, 2468, 1349, 1, 2, 2, 2, ..., 2]
        """
        src_list.append(
            pad(
                processed_src,
                (0, max_length - len(processed_src),),
                value=pad_id,
            )
        )
        tgt_list.append(
            pad(
                processed_tgt,
                (0, max_length - len(processed_tgt),),
                value=pad_id,
            )
        )

    # 将多个src句子堆叠到一起
    src = torch.stack(src_list)
    tgt = torch.stack(tgt_list)

    # tgt_y是目标句子去掉第一个token，即去掉<bos>
    tgt_y = tgt[:, 1:]
    # tgt是目标句子去掉最后一个token
    tgt = tgt[:, :-1]

    # 计算本次batch要预测的token数
    n_tokens = (tgt_y != 2).sum()

    # 返回batch后的结果
    return src, tgt, tgt_y, n_tokens
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

关于tgt和tgt_y的处理，可以参考这篇博客

有了collate_fn函数，我们就可以构造dataloader了。

train_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)
1

src, tgt, tgt_y, n_tokens = next(iter(train_loader))
src, tgt, tgt_y = src.to(device), tgt.to(device), tgt_y.to(device)
1
2

print("src.size:", src.size())
print("tgt.size:", tgt.size())
print("tgt_y.size:", tgt_y.size())
print("n_tokens:", n_tokens)
1
2
3
4

src.size: torch.Size([64, 72])
tgt.size: torch.Size([64, 71])
tgt_y.size: torch.Size([64, 71])
n_tokens: tensor(1227)
1
2
3
4

接下来，我们就可以来构建翻译模型了。

模型构建

由于nn.Transformer并没有Positional Encoding部分的实现，所以我们需要自己实现。这里我们就直接拿别人实现好的过来用：

class PositionalEncoding(nn.Module):
    "Implement the PE function."

    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)

        # 初始化Shape为(max_len, d_model)的PE (positional encoding)
        pe = torch.zeros(max_len, d_model).to(device)
        # 初始化一个tensor [[0, 1, 2, 3, ...]]
        position = torch.arange(0, max_len).unsqueeze(1)
        # 这里就是sin和cos括号中的内容，通过e和ln进行了变换
        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )
        # 计算PE(pos, 2i)
        pe[:, 0::2] = torch.sin(position * div_term)
        # 计算PE(pos, 2i+1)
        pe[:, 1::2] = torch.cos(position * div_term)
        # 为了方便计算，在最外面在unsqueeze出一个batch
        pe = pe.unsqueeze(0)
        # 如果一个参数不参与梯度下降，但又希望保存model的时候将其保存下来
        # 这个时候就可以用register_buffer
        self.register_buffer("pe", pe)

    def forward(self, x):
        """
        x 为embedding后的inputs，例如(1,7, 128)，batch size为1,7个单词，单词维度为128
        """
        # 将x和positional encoding相加。
        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32

接下来我们来定义Transformer翻译模型，nn.Transformer只实现了Transformer中下图绿色的部分，所以其他部分需要我们自己来实现：

class TranslationModel(nn.Module):

    def __init__(self, d_model, src_vocab, tgt_vocab, dropout=0.1):
        super(TranslationModel, self).__init__()

        # 定义原句子的embedding
        self.src_embedding = nn.Embedding(len(src_vocab), d_model, padding_idx=2)
        # 定义目标句子的embedding
        self.tgt_embedding = nn.Embedding(len(tgt_vocab), d_model, padding_idx=2)
        # 定义posintional encoding
        self.positional_encoding = PositionalEncoding(d_model, dropout, max_len=max_length)
        # 定义Transformer
        self.transformer = nn.Transformer(d_model, dropout=dropout, batch_first=True)

        # 定义最后的预测层，这里并没有定义Softmax，而是把他放在了模型外。
        self.predictor = nn.Linear(d_model, len(tgt_vocab))

    def forward(self, src, tgt):
        """
        进行前向传递，输出为Decoder的输出。注意，这里并没有使用self.predictor进行预测，
        因为训练和推理行为不太一样，所以放在了模型外面。
        :param src: 原batch后的句子，例如[[0, 12, 34, .., 1, 2, 2, ...], ...]
        :param tgt: 目标batch后的句子，例如[[0, 74, 56, .., 1, 2, 2, ...], ...]
        :return: Transformer的输出，或者说是TransformerDecoder的输出。
        """

        """
        生成tgt_mask，即阶梯型的mask，例如：
        [[0., -inf, -inf, -inf, -inf],
        [0., 0., -inf, -inf, -inf],
        [0., 0., 0., -inf, -inf],
        [0., 0., 0., 0., -inf],
        [0., 0., 0., 0., 0.]]
        tgt.size()[-1]为目标句子的长度。
        """
        tgt_mask = nn.Transformer.generate_square_subsequent_mask(tgt.size()[-1]).to(device)
        # 掩盖住原句子中<pad>的部分，例如[[False,False,False,..., True,True,...], ...]
        src_key_padding_mask = TranslationModel.get_key_padding_mask(src)
        # 掩盖住目标句子中<pad>的部分
        tgt_key_padding_mask = TranslationModel.get_key_padding_mask(tgt)

        # 对src和tgt进行编码
        src = self.src_embedding(src)
        tgt = self.tgt_embedding(tgt)
        # 给src和tgt的token增加位置信息
        src = self.positional_encoding(src)
        tgt = self.positional_encoding(tgt)

        # 将准备好的数据送给transformer
        out = self.transformer(src, tgt,
                               tgt_mask=tgt_mask,
                               src_key_padding_mask=src_key_padding_mask,
                               tgt_key_padding_mask=tgt_key_padding_mask)

        """
        这里直接返回transformer的结果。因为训练和推理时的行为不一样，
        所以在该模型外再进行线性层的预测。
        """
        return out

    @staticmethod
    def get_key_padding_mask(tokens):
        """
        用于key_padding_mask
        """
        return tokens == 2

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67

在nn.Transformer中，mask的-inf表示遮掩，而0表示不遮掩。而key_padding_mask的True表示遮掩，False表示不遮掩。

if model_checkpoint:
    model = torch.load(model_dir / model_checkpoint)
else:
    model = TranslationModel(256, en_vocab, zh_vocab)
model = model.to(device)
1
2
3
4
5

尝试调用一下model，验证一下是否能正常运行

model(src, tgt).size()
1

torch.Size([64, 71, 256])
1

model(src, tgt)
1

tensor([[[ 0.3853, -0.8223,  0.5280,  ..., -2.4575,  2.5116, -0.5928],
         [ 1.5033, -0.3207,  0.5466,  ..., -2.5268,  2.2986, -1.6524],
         [ 0.7981,  0.4327,  0.5015,  ..., -2.1362,  0.7818, -1.1500],
         ...,
         [ 0.6166, -0.8814, -0.0232,  ..., -1.6519,  2.8955, -1.2634],
         [ 1.9665, -0.6462, -0.0716,  ..., -2.0842,  1.7766, -0.9148],
         [ 0.9839, -0.6833,  0.2441,  ..., -1.2677,  2.3247, -1.7913]]],
       device='cuda:0', grad_fn=<NativeLayerNormBackward0>)
1
2
3
4
5
6
7
8

模型正常调用，其中71是因为tgt去掉了最后一个token。

模型训练

简单起见，本次模型训练使用的是Adam优化器，对于学习率就不进行Warmup了。

optimizer = torch.optim.Adam(model.parameters(), lr=3e-4)
1

class TranslationLoss(nn.Module):

    def __init__(self):
        super(TranslationLoss, self).__init__()
        # 使用KLDivLoss，不需要知道里面的具体细节。
        self.criterion = nn.KLDivLoss(reduction="sum")
        self.padding_idx = 2

    def forward(self, x, target):
        """
        损失函数的前向传递
        :param x: 将Decoder的输出再经过predictor线性层之后的输出。
                  也就是Linear后、Softmax前的状态
        :param target: tgt_y。也就是label，例如[[1, 34, 15, ...], ...]
        :return: loss
        """

        """
        由于KLDivLoss的input需要对softmax做log，所以使用log_softmax。
        等价于：log(softmax(x))
        """
        x = log_softmax(x, dim=-1)

        """
        构造Label的分布，也就是将[[1, 34, 15, ...]] 转化为:
        [[[0, 1, 0, ..., 0],
          [0, ..., 1, ..,0],
          ...]],
        ...]
        """
        # 首先按照x的Shape构造出一个全是0的Tensor
        true_dist = torch.zeros(x.size()).to(device)
        # 将对应index的部分填充为1
        true_dist.scatter_(1, target.data.unsqueeze(1), 1)
        # 找出<pad>部分，对于<pad>标签，全部填充为0，没有1，避免其参与损失计算。
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            true_dist.index_fill_(0, mask.squeeze(), 0.0)

        # 计算损失
        return self.criterion(x, true_dist.clone().detach())
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41

criteria = TranslationLoss()
1

完成了损失定义，就可以正式开始训练模型了，训练过程和正常模型训练相差不大，这里我使用tensorboard来记录损失：

writer = SummaryWriter(log_dir='runs/transformer_loss')
1

你可以在当前目录下运行tensorboard --logdir runs命令来启动tensorboard。

torch.cuda.empty_cache()
1

step = 0

if model_checkpoint:
    step = int('model_10000.pt'.replace("model_", "").replace(".pt", ""))

model.train()
for epoch in range(epochs):
    loop = tqdm(enumerate(train_loader), total=len(train_loader))
    for index, data in enumerate(train_loader):
        # 生成数据
        src, tgt, tgt_y, n_tokens = data
        src, tgt, tgt_y = src.to(device), tgt.to(device), tgt_y.to(device)

        # 清空梯度
        optimizer.zero_grad()
        # 进行transformer的计算
        out = model(src, tgt)
        # 将结果送给最后的线性层进行预测
        out = model.predictor(out)

        """
        计算损失。由于训练时我们的是对所有的输出都进行预测，所以需要对out进行reshape一下。
                我们的out的Shape为(batch_size, 词数, 词典大小)，view之后变为：
                (batch_size*词数, 词典大小)。
                而在这些预测结果中，我们只需要对非<pad>部分进行，所以需要进行正则化。也就是
                除以n_tokens。
        """
        loss = criteria(out.contiguous().view(-1, out.size(-1)), tgt_y.contiguous().view(-1)) / n_tokens
        # 计算梯度
        loss.backward()
        # 更新参数
        optimizer.step()

        loop.set_description("Epoch {}/{}".format(epoch, epochs))
        loop.set_postfix(loss=loss.item())
        loop.update(1)

        step += 1

        del src
        del tgt
        del tgt_y

        if step != 0 and step % save_after_step == 0:
            torch.save(model, model_dir / f"model_{step}.pt")
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45

Epoch 0/10:  78%|███████▊  | 121671/156250 [9:17:29<2:37:46,  3.65it/s, loss=2.25]
1

模型推理

训练完模型后，我们来使用我们的模型来进行一波推理。

Transformer推理时，tgt是一次一个的将token传给Transformer，例如，首次tgt为<bos>，然后预测出I，然后第二次tgt为<bos> I，预测出like，第三次tgt为<bos> I like，以此类推，直到预测结果为<eos>，或者达到句子最大长度。

model = model.eval()
1

def translate(src: str):
    """
    :param src: 英文句子，例如 "I like machine learning."
    :return: 翻译后的句子，例如：”我喜欢机器学习“
    """

    # 将与原句子分词后，通过词典转为index，然后增加<bos>和<eos>
    src = torch.tensor([0] + en_vocab(en_tokenizer(src)) + [1]).unsqueeze(0).to(device)
    # 首次tgt为<bos>
    tgt = torch.tensor([[0]]).to(device)
    # 一个一个词预测，直到预测为<eos>，或者达到句子最大长度
    for i in range(max_length):
        # 进行transformer计算
        out = model(src, tgt)
        # 预测结果，因为只需要看最后一个词，所以取`out[:, -1]`
        predict = model.predictor(out[:, -1])
        # 找出最大值的index
        y = torch.argmax(predict, dim=1)
        # 和之前的预测结果拼接到一起
        tgt = torch.concat([tgt, y.unsqueeze(0)], dim=1)
        # 如果为<eos>，说明预测结束，跳出循环
        if y == 1:
            break
    # 将预测tokens拼起来
    tgt = ''.join(zh_vocab.lookup_tokens(tgt.squeeze().tolist())).replace("<s>", "").replace("</s>", "")
    return tgt
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

translate("Alright, this project is finished. Let's see how good this is.")
1

'好吧，这个项目完成了。让我们看看这是多好的。'
1

---