如何利用PyTorch写一个Transformer实现英德互译

数据集中每一行是一对英语，德语句子对

Transformer模型出处：2017 《Attention is all you need》

Transformer中的位置编码是什么意思？

https://kazemnejad.com/blog/transformer_architecture_positional_encoding/​kazemnejad.com

运行以下代码请确保：

PyTorch=1.9.0

torchtext=0.10.0

否则肯定报错

# 利用Transformer结构
# 实现英语-德语互译

# 分词工具
from torchtext.data.utils import get_tokenizer
# 构建数据集的词典的工具
from torchtext.vocab import build_vocab_from_iterator
# Multi30k数据集
# 常用的机器翻译数据集
from torchtext.datasets import Multi30k
from typing import Iterable, List

# 从德语翻译到英语
SRC_LANGUAGE = 'de'
TGT_LANGUAGE = 'en'
# 转换器
token_transform = {}
vocab_transform = {}
# 请确保安装了以下依赖包
# pip install -U spacy
# 定义指定的分词格式
# 按照以下网址中的安装方法进行安装
# 下载指定的分词格式
# https://blog.csdn.net/yyp1998/article/details/117292281
# get_torkenizer(第一个参数指定使用的分词器，如果没有,按照空格进行分割，language=指定使用的分词格式（哪种语言）)
# 德语分词器
token_transform[SRC_LANGUAGE] = get_tokenizer('spacy', language='de_core_news_sm')
# 英语分词器
token_transform[TGT_LANGUAGE] = get_tokenizer('spacy', language='en_core_web_sm')


def yield_tokens(data_iter: Iterable, language: str) -> List[str]:
    # 利用上面的分词器
    # 对数据集中的每行英语<->德语对进行分词处理
    # 标记是英语还是德语
    language_index = {SRC_LANGUAGE: 0, TGT_LANGUAGE: 1}
    # 迭代数据集中的每一行
    for data_sample in data_iter:
        # 产生这一行的分词
        yield token_transform[language](data_sample[language_index[language]])


# 定义特殊标记
# UNK=unknown 未知词语标记
# PAD=padding 序列填充标记
# BOS=begin of string(有的也定义为SOS，start of string) 一句话的开始标记
# EOS=end of string 一句话的结尾标记
UNK_IDX, PAD_IDX, BOS_IDX, EOS_IDX = 0, 1, 2, 3
special_symbols = ['', '', '', '']
for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:
    # split指定只获取训练数据集
    # language_pair指定我们想要使用的句子对
    train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))
    # 根据数据集建立一个词典
    # 加入我们上面定义的特殊标记
    # 并把这些特殊标记放在词典的最前面
    # 最小词频指定为1，意思是只要训练集中出现过的词语，通通加入到词典中
    vocab_transform[ln] = build_vocab_from_iterator(yield_tokens(train_iter, ln),
                                                    min_freq=1,
                                                    specials=special_symbols,
                                                    special_first=True)

# 配置UNK_IDX为默认索引
# 后面翻译过程中要是出现词典中没见过的词语，一律判定为UNK
for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:
    vocab_transform[ln].set_default_index(UNK_IDX)
# 测试
# 获取am begin在我们建立的英语词典中的索引
print(vocab_transform['en'](['am', 'begin']))

from torch import Tensor
import torch
import torch.nn as nn
from torch.nn import Transformer
import math

# 获取可用设备
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')


# 位置编码
# 详细的看我的上几篇中的《如何利用PyTorch训练一个Transformer语言模型学习词嵌入》
class PositionalEncoding(nn.Module):
    def __init__(self, emb_size: int, dropout: float, maxlen: int = 5000):
        '''
        :param emb_size: 词嵌入的维度
        :param dropout: 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))
        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):
        '''
        :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)


# Transformer
class Seq2SeqTransformer(nn.Module):
    def __init__(self, num_encoder_layers: int, num_decoder_layers: int, emb_size: int, nhead: 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 nhead: 多头注意力头数
        :param src_vocab_size: 源语言词典大小
        :param tgt_vocab_size: 目标语言词典大小
        :param dim_feedforward: 隐藏层的维度大小
        :param dropout: dropout概率
        '''
        super(Seq2SeqTransformer, self).__init__()
        self.transformer = Transformer(d_model=emb_size,
                                       nhead=nhead,
                                       num_encoder_layers=num_encoder_layers,
                                       num_decoder_layers=num_decoder_layers,
                                       dim_feedforward=dim_feedforward,
                                       dropout=dropout)
        # 输出层
        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: 对于输入的Mask矩阵
        :param tgt_mask: 对于输出的Mask矩阵
        :param src_padding_mask:对于填充的Mask矩阵
        :param tgt_padding_mask:
        :param memory_key_padding_mask:Encoder输出的Mask矩阵
        :return:
        '''
        # 先进行词嵌入，然后进行位置编码
        src_emb = self.positional_encoding(self.src_tok_emb(src))
        tgt_emb = self.positional_encoding(self.tgt_tok_emb(trg))
        # 输入到transformer中进行编、解码
        outs = self.transformer(src_emb, tgt_emb, src_mask, tgt_mask, None,
                                src_padding_mask, tgt_padding_mask, memory_key_padding_mask)
        # 产生transformer的输出
        return self.generator(outs)

    # Encoder
    def encode(self, src: Tensor, src_mask: Tensor):
        return self.transformer.encoder(self.positional_encoding(
            self.src_tok_emb(src)), src_mask)

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


def generate_square_subsequent_mask(sz):
    # torch.ones((sz, sz), device=DEVICE)生成一个指定设备上元素都为1的[sz,sz]矩阵
    # torch.triu()函数用于从矩阵中沿着主对角线截取一部分元素，diagonal=参数控制截取的位置
    # 如果不指定默认diagonal=0,就是截取一个矩阵的上三角部分
    # torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1转换为一个bool值的mask矩阵
    # .transpose(0,1)是矩阵转置
    mask = (torch.triu(torch.ones((sz, sz), device=DEVICE)) == 1).transpose(0, 1)
    # 将False元素设置为-inf
    # 其余元素设置为0.0
    mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
    # 上面的内容可以写为
    # mask=torch.triu(torch.ones(sz, sz) * float('-inf'), diagonal=1)
    return mask


print(generate_square_subsequent_mask(3))


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)
    # 如果指定了Padding，进行全0填充
    # 词嵌入中的padding是做什么的呢
    # 为了方便模型的训练
    # 如果某句话的长度没有达到这批次数据每句话的最大长度，进行padding填充对齐
    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


torch.manual_seed(0)
# 词典的大小
SRC_VOCAB_SIZE = len(vocab_transform[SRC_LANGUAGE])
TGT_VOCAB_SIZE = len(vocab_transform[TGT_LANGUAGE])
# 词嵌入维度
EMB_SIZE = 512
# 多头注意力头数
NHEAD = 8
# Transformer隐藏层的维度
FFN_HID_DIM = 512
# 批大小
BATCH_SIZE = 128
# 编、解码器的层数
NUM_ENCODER_LAYERS = 3
NUM_DECODER_LAYERS = 3
# 实例化一个Transformer
transformer = Seq2SeqTransformer(NUM_ENCODER_LAYERS, NUM_DECODER_LAYERS, EMB_SIZE,
                                 NHEAD, SRC_VOCAB_SIZE, TGT_VOCAB_SIZE, FFN_HID_DIM)
# 权重初始化
for p in transformer.parameters():
    if p.dim() > 1:
        nn.init.xavier_uniform_(p)
# 如果GPU可用放在GPU上
transformer = transformer.to(DEVICE)
# 损失函数
# padding的不是真正的词语，要忽略掉 ignore_index=PAD_IDX
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)
from torch.nn.utils.rnn import pad_sequence


# 转换器聚合器
# 组合多个转换器对输入进行转换
# 类似torchvision中的Compose
# 它的原理是什么呢
# 首先*transforms这个动作表示解包操作，拆分输入进来的多个转换器
# 然后通过内置函数依次利用不同的转换器处理输入
# 相当于一个递归，但是每次调用func时使用的转换器不同
def sequential_transforms(*transforms):
    def func(txt_input):
        for transform in transforms:
            txt_input = transform(txt_input)
        return txt_input

    return func


# 在一句话的开头和结尾添加特殊标记
def tensor_transform(token_ids: List[int]):
    return torch.cat((torch.tensor([BOS_IDX]),
                      torch.tensor(token_ids),
                      torch.tensor([EOS_IDX])))


# 组合转换器
text_transform = {}
for ln in [SRC_LANGUAGE, TGT_LANGUAGE]:
    # 先分词
    # 然后获取每个词语在词典中的索引
    # 然后转换为张量
    text_transform[ln] = sequential_transforms(token_transform[ln],
                                               vocab_transform[ln],
                                               tensor_transform)


# 处理一批次的数据
def collate_fn(batch):
    src_batch, tgt_batch = [], []
    for src_sample, tgt_sample in batch:
        # rstrip()删除字符串末尾的指定字符
        # 这两行在做什么呢
        # 根据我们前面组合的转换器
        # 对于一句话
        # 先分词
        # 然后获取这句话每个词在词典中的索引
        # 然后加入开始，结束标记，转换为张量
        # 然后加入到这批次数据中
        src_batch.append(text_transform[SRC_LANGUAGE](src_sample.rstrip("\n")))
        tgt_batch.append(text_transform[TGT_LANGUAGE](tgt_sample.rstrip("\n")))
    # 拿src_batch举例
    # 处理完后src_batch是一个长度为batch_size的列表，其中每个元素是[seq_len,源语言词典大小]的矩阵
    # 但是注意这里seq_len长度不一定都相同
    # 因为你无法保证每一句话的长度都是相同的
    # 所以利用pad_sequence进行填充，填充到相同的长度
    # 方便处理
    # 但是这些填充又不是真正的词语
    # 所以我们要引入Padding Mask矩阵遮盖住这些地方，防止对翻译产生影响
    src_batch = pad_sequence(src_batch, padding_value=PAD_IDX)
    tgt_batch = pad_sequence(tgt_batch, padding_value=PAD_IDX)
    # pad_sequence之后
    # src_batch的维度为[统一的seq_len,batch_size,源语言词典大小]
    return src_batch, tgt_batch


from torch.utils.data import DataLoader


def train_epoch(model, optimizer):
    # 开启模型训练模式
    model.train()
    losses = 0
    # 获取一批次的数据
    # 并指定对于这批数据使用我们上面定义的处理方法， 转换为张量
    train_iter = Multi30k(split='train', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))
    train_dataloader = DataLoader(train_iter, batch_size=BATCH_SIZE, collate_fn=collate_fn)
    # 训练
    for src, tgt in train_dataloader:
        # 如果GPU可用放在GPU上
        src = src.to(DEVICE)
        tgt = tgt.to(DEVICE)
        # 最后一个是结尾特殊标志
        # 省去
        tgt_input = tgt[:-1, :]
        # 获取Mask矩阵
        src_mask, tgt_mask, src_padding_mask, tgt_padding_mask = create_mask(src, tgt_input)
        # 获取模型的输出
        # logits是没有经过softmax的模型输出的意思
        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:, :]
        # logits的维度是[seq_len-1,batch_size,目标语言的词典大小]
        # logits.reshape(-1, logits.shape[-1])后维度变为[(seq_len-1)*batch_size,目标语言的词典大小]
        # tgt_out维度是[seq_len-1,batch_size,目标语言的词典大小]
        # tgt_out.reshape(-1)后tgt_out的维度是[seq_len-1*batch_size*目标语言的词典大小]
        # 为什么要这么做呢
        # 可以类比图片分类
        # tgt_out的最后一个维度是真实标签
        # 而logits的最后一个维度是对于每一个词语的预测
        loss = loss_fn(logits.reshape(-1, logits.shape[-1]), tgt_out.reshape(-1))
        # 反向传播
        loss.backward()
        # 梯度更新
        optimizer.step()
        losses += loss.item()
    return losses / len(train_dataloader)


# 模型验证
def evaluate(model):
    model.eval()
    losses = 0
    val_iter = Multi30k(split='valid', language_pair=(SRC_LANGUAGE, TGT_LANGUAGE))
    val_dataloader = DataLoader(val_iter, batch_size=BATCH_SIZE, collate_fn=collate_fn)
    for src, tgt in val_dataloader:
        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_dataloader)


from timeit import default_timer as timer

NUM_EPOCHS = 18
# 迭代训练和验证
for epoch in range(1, NUM_EPOCHS + 1):
    start_time = timer()
    train_loss = train_epoch(transformer, optimizer)
    end_time = timer()
    val_loss = evaluate(transformer)
    print((
        f"Epoch: {epoch}, Train loss: {train_loss:.3f}, Val loss: {val_loss:.3f}, "f"Epoch time = {(end_time - start_time):.3f}s"))


# 让训练完的模型产生翻译字符串
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)
        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: torch.nn.Module, src_sentence: str):
    model.eval()
    src = text_transform[SRC_LANGUAGE](src_sentence).view(-1, 1)
    num_tokens = src.shape[0]
    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(vocab_transform[TGT_LANGUAGE].lookup_tokens(list(tgt_tokens.cpu().numpy()))).replace("",
                                                                                                         "").replace(
        "", "")


# 将德语翻译为英语
print(translate(transformer, "Eine Gruppe von Menschen steht vor einem Iglu ."))