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

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

本篇最终效果：

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

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

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

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

torch>=1.11.0
tokenizers==0.12.1
torchtext==0.12.0
tensorboard==2.8.0

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

import os
import math

import torch
import torch.nn as nn

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

下载数据集。数据集包括两个文件，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")

model_checkpoint = None

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)

en_vocab = None
zh_vocab = None

batch_size = 64

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)

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

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

文本分词与构造词典

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

接下来来构造英文词典：

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', '.']
"""

    return tokenizer.encode(line, add_special_tokens=False).tokens

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

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

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

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

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

上面的分词结果中，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()

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

else:

    en_vocab = build_vocab_from_iterator(

        yield_en_tokens(),

        min_freq=2,

        specials=["", "", "", ""],
    )

    en_vocab.set_default_index(en_vocab[""])

    if use_cache:
        torch.save(en_vocab, en_vocab_file)

print("英文词典大小:", len(en_vocab))
print(dict((i, en_vocab.lookup_token(i)) for i in range(10)))

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

接着我们来构建中文词典，中文词比较多，容易产生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()

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=["", "", "", ""],
    )
    zh_vocab.set_default_index(zh_vocab[""])
    torch.save(zh_vocab, zh_vocab_file)

print("中文词典大小:", len(zh_vocab))
print(dict((i, zh_vocab.lookup_token(i)) for i in range(10)))

中文词典大小: 8280
{0: '<s>', 1: '</s>', 2: '<pad>', 3: '<unk>', 4: '&#x3002;', 5: '&#x7684;', 6: '&#xFF0C;', 7: '&#x6211;', 8: '&#x4F60;', 9: '&#x662F;'}
</unk></pad>

Dataset and Dataloader

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

class TranslationDataset(Dataset):

    def __init__(self):

        self.en_tokens = self.load_tokens(en_filepath, en_tokenizer, en_vocab, "构建英文tokens", 'en')

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

        tokens_list = []

        with open(file, encoding='utf-8') as file:

            for line in tqdm(file, desc=desc, total=row_count):

                tokens = tokenizer(line)

                tokens = vocab(tokens)

                tokens_list.append(tokens)

        if use_cache:
            torch.save(tokens_list, cache_file)

        return tokens_list

dataset = TranslationDataset()

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

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

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

print(dataset.__getitem__(0))

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

Dataset中并不包含 <bos></bos>和 <eos></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，即
             n_tokens：tgt_y中的token数，不计算在内。
"""

    bs_id = torch.tensor([0])

    eos_id = torch.tensor([1])

    pad_id = 2

    src_list, tgt_list = [], []

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

        processed_src = torch.cat(

            [
                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 = torch.stack(src_list)
    tgt = torch.stack(tgt_list)

    tgt_y = tgt[:, 1:]

    tgt = tgt[:, :-1]

    n_tokens = (tgt_y != 2).sum()

    return src, tgt, tgt_y, n_tokens

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

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

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

print("src.size:", src.size())
print("tgt.size:", tgt.size())
print("tgt_y.size:", tgt_y.size())
print("n_tokens:", n_tokens)

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

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

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

        pe = torch.zeros(max_len, d_model).to(device)

        position = torch.arange(0, max_len).unsqueeze(1)

        div_term = torch.exp(
            torch.arange(0, d_model, 2) * -(math.log(10000.0) / d_model)
        )

        pe[:, 0::2] = torch.sin(position * div_term)

        pe[:, 1::2] = torch.cos(position * div_term)

        pe = pe.unsqueeze(0)

        self.register_buffer("pe", pe)

    def forward(self, x):
"""
        x 为embedding后的inputs，例如(1,7, 128)，batch size为1,7个单词，单词维度为128
"""

        x = x + self.pe[:, : x.size(1)].requires_grad_(False)
        return self.dropout(x)

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

class TranslationModel(nn.Module):

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

        self.src_embedding = nn.Embedding(len(src_vocab), d_model, padding_idx=2)

        self.tgt_embedding = nn.Embedding(len(tgt_vocab), d_model, padding_idx=2)

        self.positional_encoding = PositionalEncoding(d_model, dropout, max_len=max_length)

        self.transformer = nn.Transformer(d_model, dropout=dropout, batch_first=True)

        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)

        src_key_padding_mask = TranslationModel.get_key_padding_mask(src)

        tgt_key_padding_mask = TranslationModel.get_key_padding_mask(tgt)

        src = self.src_embedding(src)
        tgt = self.tgt_embedding(tgt)

        src = self.positional_encoding(src)
        tgt = self.positional_encoding(tgt)

        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

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

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

model(src, tgt).size()

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

model(src, tgt)

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

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

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

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

class TranslationLoss(nn.Module):

    def __init__(self):
        super(TranslationLoss, self).__init__()

        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],
          ...]],
        ...]
"""

        true_dist = torch.zeros(x.size()).to(device)

        true_dist.scatter_(1, target.data.unsqueeze(1), 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())

criteria = TranslationLoss()

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

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

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

torch.cuda.empty_cache()

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

        out = model(src, tgt)

        out = model.predictor(out)

"""
        计算损失。由于训练时我们的是对所有的输出都进行预测，所以需要对out进行reshape一下。
                我们的out的Shape为(batch_size, 词数, 词典大小)，view之后变为：
                (batch_size*词数, 词典大小)。
                而在这些预测结果中，我们只需要对非部分进行，所以需要进行正则化。也就是
                除以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")

Epoch 0/10:  78%|&#x2588;&#x2588;&#x2588;&#x2588;&#x2588;&#x2588;&#x2588;&#x258A;  | 121671/156250 [9:17:29<2:37:46, 3.65it s, loss="2.25]" < code></2:37:46,>

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

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

model = model.eval()

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

    src = torch.tensor([0] + en_vocab(en_tokenizer(src)) + [1]).unsqueeze(0).to(device)

    tgt = torch.tensor([[0]]).to(device)

    for i in range(max_length):

        out = model(src, tgt)

        predict = model.predictor(out[:, -1])

        y = torch.argmax(predict, dim=1)

        tgt = torch.concat([tgt, y.unsqueeze(0)], dim=1)

        if y == 1:
            break

    tgt = ''.join(zh_vocab.lookup_tokens(tgt.squeeze().tolist())).replace("", "").replace("", "")
    return tgt

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

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

Original: https://blog.csdn.net/zhaohongfei_358/article/details/126175328
Author: iioSnail
Title: Pytorch入门实战（5）：基于nn.Transformer实现机器翻译（英译汉）