Transformer pytorch实现
目录：
- 前言
- 数据处理
- Transformer各个模块具体实现
- 词嵌入层
- 位置编码(positional encoding）
- 编码器
- 多头自注意力
- 层归一化
- 残差连接
- 逐位置前馈网络(Position-wise Feed-Forward Networks）
- 编码器整体架构
- 解码器
- Transformer模型整体架构
- 模型训练及预测
- 标签平滑
- 计算损失
- 优化器
- 训练
- 预测
前言
我们在前面介绍了Transformer的理论，但是始终只是纸上谈兵，本文我们利用PyTorch这个深度学习库，来具体的实现一下，主要参考的是The Annotated Transformer，并完成一个简单的中英文机器翻译任务。
在正式介绍之前，我们先做一些准备工作，首先是导入一些包和初始化工作：
import copy import math import matplotlib.pyplot as plt import numpy as np import os import seaborn as sns import time import torch import torch.nn as nn import torch.nn.functional as F from collections import Counter #简繁体转化包 from langconv import Converter from nltk import word_tokenize from torch.autograd import Variable
下载NLTK后，我们进行以下设置：
import nltk
nltk.set_proxy('http://proxy.example.com:3128', ('USERNAME', 'PASSWORD'))
nltk.download()
初始化参数设置：
# 初始化参数设置
PAD = 0 # padding占位符的索引
UNK = 1 # 未登录词标识符的索引
BATCH_SIZE = 128 # 批次大小
EPOCHS = 20 # 训练轮数
LAYERS = 6 # transformer中encoder、decoder层数
H_NUM = 8 # 多头注意力个数
D_MODEL = 256 # 输入、输出词向量维数
D_FF = 1024 # feed forward全连接层维数
DROPOUT = 0.1 # dropout比例
MAX_LENGTH = 60 # 语句最大长度
TRAIN_FILE = 'nmt/en-cn/train.txt' # 训练集
DEV_FILE = "nmt/en-cn/dev.txt" # 验证集
SAVE_FILE = 'save/model.pt' # 模型保存路径
DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
数据处理
数据集可以到这里来下载，里面有这几个数据文件，
训练数据是长这样子的：
Anyone can do that. 任何人都可以做到。 How about another piece of cake? 要不要再來一塊蛋糕？ She married him. 她嫁给了他。 I don't like learning irregular verbs. 我不喜欢学习不规则动词。 It's a whole new ball game for me. 這對我來說是個全新的球類遊戲。 He's sleeping like a baby. 他正睡着，像个婴儿一样。 He can play both tennis and baseball. 他既会打网球，又会打棒球。 We should cancel the hike. 我們應該取消這次遠足。 He is good at dealing with children. 他擅長應付小孩子。
对于数据预处理，主要做的事是：
- 加载数据，
- 分词，
- 构建词表，
- 批次划分，
对中文语句一般以字为单位进行切分，所以不需要对中文语句分词。
首先我们按批次对数据进行填充，使得同一批次内的数据长度对齐，不同批次内的最大长度可以不一致。
def seq_padding(X, padding=PAD):
"""
按批次（batch）对数据填充、长度对齐
"""
# 计算该批次各条样本语句长度
lens = [len(x) for x in X]
# 获取该批次样本中语句长度最大值
max_len = max(lens)
# 遍历该批次样本，如果语句长度小于最大长度，则用padding填充
return np.array([
np.concatenate([x, [padding] * (max_len - len(x))]) if len(x) < max_len else x for x in X
])
如上所见，这里的中文为繁体，所以我们还需要做简繁转换：
def cht_to_chs(sent):
sent = Converter("zh-hans").convert(sent)
sent.encode("utf-8")
return sent
下面来准备数据：
class PrepareData:
def __init__(self, train_file, dev_file):
# 读取数据、分词
self.train_en, self.train_cn = self.load_data(train_file)
self.dev_en, self.dev_cn = self.load_data(dev_file)
# 构建词表
self.en_word_dict, self.en_total_words, self.en_index_dict = \
self.build_dict(self.train_en)
self.cn_word_dict, self.cn_total_words, self.cn_index_dict = \
self.build_dict(self.train_cn)
# 单词映射为索引
self.train_en, self.train_cn = self.word2id(self.train_en, self.train_cn, self.en_word_dict, self.cn_word_dict)
self.dev_en, self.dev_cn = self.word2id(self.dev_en, self.dev_cn, self.en_word_dict, self.cn_word_dict)
# 划分批次、填充、掩码
self.train_data = self.split_batch(self.train_en, self.train_cn, BATCH_SIZE)
self.dev_data = self.split_batch(self.dev_en, self.dev_cn, BATCH_SIZE)
def load_data(self, path):
"""
读取英文、中文数据
对每条样本分词并构建包含起始符和终止符的单词列表
形式如：en = [['BOS', 'i', 'love', 'you', 'EOS'], ['BOS', 'me', 'too', 'EOS'], ...]
cn = [['BOS', '我', '爱', '你', 'EOS'], ['BOS', '我', '也', '是', 'EOS'], ...]
"""
en = []
cn = []
with open(path, mode="r", encoding="utf-8") as f:
for line in f.readlines():
sent_en, sent_cn = line.strip().split("\t")
sent_en = sent_en.lower()
sent_cn = cht_to_chs(sent_cn)
sent_en = ["BOS"] + word_tokenize(sent_en) + ["EOS"]
# 中文按字符切分
sent_cn = ["BOS"] + [char for char in sent_cn] + ["EOS"]
en.append(sent_en)
cn.append(sent_cn)
return en, cn
def build_dict(self, sentences, max_words=5e4):
"""
构造分词后的列表数据
构建单词-索引映射（key为单词，value为id值）
"""
# 统计数据集中单词词频
word_count = Counter([word for sent in sentences for word in sent])
# 按词频保留前max_words个单词构建词典
# 添加UNK和PAD两个单词
ls = word_count.most_common(int(max_words))
total_words = len(ls) + 2
word_dict = {w[0]: index + 2 for index, w in enumerate(ls)}
word_dict['PAD'] = PAD # PAD = 0，padding占位符的索引
word_dict['UNK'] = UNK # UNK = 1，未登录词标识符的索引
# 构建id2word映射
index_dict = {v: k for k, v in word_dict.items()}
return word_dict, total_words, index_dict
def word2id(self, en, cn, en_dict, cn_dict, sort=True):
"""
将英文、中文单词列表转为单词索引列表
`sort=True`表示以英文语句长度排序，以便按批次填充时，同批次语句填充尽量少
"""
length = len(en)
# 单词映射为索引
out_en_ids = [[en_dict.get(word, UNK) for word in sent] for sent in en]
out_cn_ids = [[cn_dict.get(word, UNK) for word in sent] for sent in cn]
# 按照语句长度排序，使批次内的长度尽可能一致
def len_argsort(seq):
"""
传入一系列语句数据(分好词的列表形式)，
按照语句长度排序后，返回排序后原来各语句在数据中的索引下标
"""
return sorted(range(len(seq)), key=lambda x: len(seq[x]))
# 按相同顺序对中文、英文样本排序
if sort:
# 以英文语句长度排序
sorted_index = len_argsort(out_en_ids)
out_en_ids = [out_en_ids[idx] for idx in sorted_index]
out_cn_ids = [out_cn_ids[idx] for idx in sorted_index]
return out_en_ids, out_cn_ids
def split_batch(self, en, cn, batch_size, shuffle=True):
"""
划分批次
`shuffle=True`表示对各批次顺序随机打乱
"""
# 每隔batch_size取一个索引作为后续batch的起始索引
idx_list = np.arange(0, len(en), batch_size)
# 起始索引随机打乱
if shuffle:
np.random.shuffle(idx_list)
# 存放所有批次的语句索引
batch_indexs = []
for idx in idx_list:
"""
形如[array([4, 5, 6, 7]),
array([0, 1, 2, 3]),
array([8, 9, 10, 11]),
...]
"""
# 起始索引最大的批次可能发生越界，要限定其索引
batch_indexs.append(np.arange(idx, min(idx + batch_size, len(en))))
# 构建批次列表
batches = []
for batch_index in batch_indexs:
# 按当前批次的样本索引采样
batch_en = [en[index] for index in batch_index]
batch_cn = [cn[index] for index in batch_index]
# 对当前批次中所有语句填充、对齐长度
# 维度为：batch_size * 当前批次中语句的最大长度
batch_cn = seq_padding(batch_cn)
batch_en = seq_padding(batch_en)
# 将当前批次添加到批次列表
# Batch类用于实现注意力掩码
batches.append(Batch(batch_en, batch_cn))
return batches
上面用到的批次类的实现如下：
class Batch:
"""Object for holding a batch of data with mask during training."""
"""
批次类
1. 输入序列（源）
2. 输出序列（目标）
3. 构造掩码
"""
def __init__(self, src, tgt=None, pad=PAD):
'''
:param src: 源数据 [batch_size, input_len]
:param tgt: 目标数据 [batch_size, input_len]
'''
# 将输入、输出单词id表示的数据规范成整数类型
src = torch.from_numpy(src).to(DEVICE).long()
tgt = torch.from_numpy(tgt).to(DEVICE).long()
self.src = src
# 对于当前输入的语句非空部分进行判断，bool序列
# 并在seq length前面增加一维，形成维度为 1×seq length 的矩阵
self.src_mask = (src != pad).unsqueeze(-2)
# 如果输出目标不为空，则需要对解码器使用的目标语句进行掩码
if tgt is not None:
# 解码器使用的目标输入部分
# 去掉最后一个，可用于teacher forcing，此时self.trg维度变成了[batch_size, input_len-1]
self.tgt = tgt[:, :-1]
# 去掉第一个，构建目标输出
# self.tgt_y：[batch_size,target_len-1]
self.tgt_y = tgt[:, 1:]
# 目标mask的目的是防止当前位置注意到后面的位置 [batch_size,input_len-1, input_len-1]
self.tgt_mask = self.make_std_mask(self.tgt, pad)
# 实际单词(不包括填充词)数量
self.ntokens = (self.tgt_y != pad).data.sum()
# 掩码操作
@staticmethod
def make_std_mask(tgt, pad):
"Create a mask to hide padding and future words."
# tat_mask:[batch_size,1,target_len-1]
tgt_mask = (tgt != pad).unsqueeze(-2) # 在倒数第二个位置加入一个维度
# type_as 把调用tensor的类型变成给定tensor的
tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)
return tgt_mask
数据准备好了，下面我们正式开始一步步实现Transformer模型。
Transformer各个模块具体实现
词嵌入层
可以随机初始化进行学习，也可以加载预训练的词向量，它的实现比较简单：
class Embeddings(nn.Module):
def __init__(self, d_model, vocab):
super(Embeddings, self).__init__()
# Embedding层
self.embed = nn.Embedding(vocab, d_model)
# Embedding维数
self.d_model = d_model
def forward(self, x):
# 返回x的词向量（需要乘以math.sqrt(d_model)）
return self.embed(x) * math.sqrt(self.d_model)
位置编码(positional encoding）
位置编码向量与词向量维度相同，$\text {max_seq_len} \times \text {embedding_dim}$，Transformer原文中使用正、余弦函数的线性变换对单词位置编码：
为了使用序列顺序信息，作者提出了利用不同频率的正弦和余弦函数表示位置编码。我们的位置编码应该满足下面的要求：
- 对于每个时间步(句子中的单词位置)，它都能输出独一无二的编码
- 任意两个时间步之间的距离都应该是一个常量，而不因句子长度而变
- 我们的模型应该能轻易地泛化到更长的句子，它的值应该是有界的
- 位置编码必须是确定的(个人：也就是是确定的，不用通过学习而得到)
作者提出的编码方式是一个简单且天才的技术，满足了上面所有的要求。
- 首先，它不是一个标量，而是一个包含特定位置信息的d维向量。
- 其次，该编码并没有整合到模型中。相反，这个向量用于为每个单词设置关于它在句子中位置的信息。换言之，通过注入单词的顺序来增强模型的输入。
令$t$为输入序列中某个位置，$\vec p_t$是该位置的位置编码，d是向量维度。f是通过以下公式产生位置编码向量的函数：
其中，
由该式子可以看出，频率(单位时间内往复振动的次数)是随着向量维度降低的(由$\frac{1}{2\pi}$ 降低成$\frac{1}{10000 \cdot 2\pi}$)。因此波长(也就是周期，往复振动一次所需要的时间T=2π/ω)形成一个从$2 \pi$到$10000 \cdot 2\pi$的等比数列。
下面补充一下波长和频率的计算：
我们也能想象位置编码$\vec{p_t} $ 是一个包含各个频率的正弦和余弦向量，这里d可以被2整除。
为什么正弦和余弦的组合可以表示顺序？假设我们用二进制来表示数字。
可以看到，随着十进制数的增加，每个位的变化率是不一样的，越低位的变化越快，红色位0和1，每个数字都会变化一次； 而黄色位，每8个数字才会变化一次。 但是二进制值的0 ,1是离散的，浪费了它们之间无限的浮点数。所以我们使用它们的连续浮动版本-正弦函数。 此外，通过降低它们的频率，我们可以从红色位变成黄色位，这样就实现了这种低位到高位的变换。如下图所示：
原文中提到，对于任何固定的偏移量$k，PE_{pos + k}$ 都要能表示为$PE_{pos}$的一个线性函数。对每个频率$\omega_k$相应的正-余弦对，存在一个线性转换$M \in \mathbb{R}^{2\times2}$ ：
从上可以看出，
最终的转换与$t$无关。类似地，我们可以找到其他正-余弦对的$M$，最终允许我们表示$\vec {p_{t + \phi}}$为一个$\vec {p_t}$对任意固定偏移量$\phi$的线性函数。这个属性，使模型很容易学得相对位置信息。这解释了为什么要选择交替的正弦和余弦函数，仅通过正弦或余弦函数达不到这一点。
我们实现位置编码如下：
class PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout, max_len=5000):
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
# 位置编码矩阵，维度[max_len, embedding_dim]
pe = torch.zeros(max_len, d_model, device=DEVICE)
# 单词位置
position = torch.arange(0.0, max_len, device=DEVICE)
position.unsqueeze_(1)# [max_len, 1]
# 使用exp和log实现幂运算
div_term = torch.exp(torch.arange(0.0, d_model, 2, device=DEVICE) * (- math.log(1e4) / d_model))
div_term.unsqueeze_(0)# [1,ceil(d_model/2)]
# 计算单词位置沿词向量维度的纹理值
# 这里的实现可能有bug,如果d_model是偶数是正确的，但是如果d_model是奇数则不正确
pe[:, 0 : : 2] = torch.sin(torch.mm(position, div_term))
pe[:, 1 : : 2] = torch.cos(torch.mm(position, div_term))
# 增加批次维度，[1, max_len, embedding_dim]
pe.unsqueeze_(0)
# 将位置编码矩阵注册为buffer(不参加训练)，因为它是绝对位置编码
self.register_buffer('pe', pe)
def forward(self, x):
# 将一个批次中语句所有词向量与位置编码相加
# 注意，位置编码不参与训练，因此设置requires_grad=False
x = x + self.pe[:, : x.size(1)].requires_grad_(False)
return self.dropout(x)
编码器
多头自注意力
开头的多头注意力的输入为：
$X_{embedding}$为维度为[batch_size,seq_len,embed_dim]，
这里我们假设我们的head的数量为超参数$h$，并且$h$一定得能够被$embed\_dim$整除，也就说，对于单个的head,单头注意力的输入嵌入维度为$embed\_dim/h$。为了利用张量运算，以一种紧凑的、统一的方式，一次性得出多个head所对应的多个$Q,K,V$矩阵，我们对$X_{embedding}$做线性映射，分别乘上三个权重矩阵$W_Q,W_K,W_V \in \Bbb R^{embed\_dim \times embed\_dim}$：
此时得到的$Q,K,V$的维度都为$[batch\_size,seq\_len,embed\_dim]$，为了后续进行多头自注意力操作的方便，我们将得到的$Q,K,V$矩阵进行reshape操作，将他们的维度由$[batch\_size,seq\_len,h \cdot embed\_dim/h]$，变为：$[batch\_size,seq\_len,h, embed\_dim/h]$
之后我们把$Q,K,V$中的$seq\_{len},h$进行了一下转置，转置后的$Q,K,V$的维度为$[batch\_size,h,seq\_len,embed\_dim / h]$。
对于单头自注意力的计算，可按照下面的公式进行：
$\sqrt{d_k}$式为了把注意力矩阵变成标准正态分布，使得softmax归一化之后的结果更加稳定。
下面我们来实现多头注意力，首先实现克隆帮助函数：
def clones(module, N):
"""
克隆基本单元，克隆的单元之间参数不共享
"""
return nn.ModuleList([
copy.deepcopy(module) for _ in range(N)
])
然后实现缩放注意力计算函数：
def attention(query, key, value, mask=None, dropout=None):
"""
Scaled Dot-Product Attention
"""
# q、k、v向量长度为d_k
d_k = query.size(-1)
# 矩阵乘法实现q、k点积注意力，sqrt(d_k)归一化
scores = torch.matmul(query, key.transpose(-2, -1)) / math.sqrt(d_k)
# 注意力掩码机制
if mask is not None:
# scores: [batch_size,h,seq_len,seq_len]
scores = scores.masked_fill(mask==0, -1e9)
# 注意力矩阵softmax归一化
p_attn = F.softmax(scores, dim=-1)
# dropout
if dropout is not None:
p_attn = dropout(p_attn)
# 注意力对v加权
return torch.matmul(p_attn, value), p_attn
最后来实现多头注意力层：
class MultiHeadedAttention(nn.Module):
"""
Multi-Head Attention（前面图中编码器第2部分）
"""
def __init__(self, h, d_model, dropout=0.1):
super(MultiHeadedAttention, self).__init__()
"""
`h`：注意力头的数量
`d_model`：词向量维数
"""
# 确保整除
assert d_model % h == 0
# q、k、v向量维数
self.d_k = d_model // h
# 头的数量
self.h = h
# WQ、WK、WV矩阵及多头注意力拼接变换矩阵WO
self.linears = clones(nn.Linear(d_model, d_model), 4)
self.attn = None
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):# mask:[batch_size,seq_len]
if mask is not None:
mask = mask.unsqueeze(1)# [batch_size,1,seq_len]
# 批次大小
nbatches = query.size(0)
# WQ、WK、WV分别对词向量线性变换，并将结果拆成h块
query, key, value = [
l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
for l, x in zip(self.linears, (query, key, value))
]
# 注意力加权
x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
# 多头注意力加权拼接
x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.h * self.d_k)
# 对多头注意力加权拼接结果线性变换
return self.linears[-1](x)
层归一化
层归一化针对每个输入的每个维度进行归一化操作。假设有$H$个维度，$x=(x_1,x_2,\cdots,x_H)$，层归一化首先计算这$H$个维度的均值和方差，然后进行归一化得到$N(x)$，接着做一个缩放.
$h$就是LN层的输出；$\mu$和$\sigma$是输入各个维度的均值和方差；$\alpha$和$\beta$是两个可学习的参数；和$h$ 的维度相同。
class LayerNorm(nn.Module):
def __init__(self, features, eps=1e-6):
super(LayerNorm, self).__init__()
# α、β分别初始化为1、0
self.a_2 = nn.Parameter(torch.ones(features))
self.b_2 = nn.Parameter(torch.zeros(features))
# 平滑项
self.eps = eps
def forward(self, x):# x:[batch_size,seq_len,embedding_dim]
# 沿词向量方向计算均值和方差
# mean:[batch_size,seq_len,1]
mean = x.mean(dim=-1, keepdim=True)
# std:[batch_size,seq_len,1]
std = x.std(dim=-1, keepdim=True)
# 沿词向量和语句序列方向计算均值和方差
# mean = x.mean(dim=[-2, -1], keepdim=True)
# std = x.std(dim=[-2, -1], keepdim=True)
# 归一化
x = (x - mean) / torch.sqrt(std ** 2 + self.eps)
return self.a_2 * x + self.b_2
残差连接
That is, the output of each sub-layer is $LayerNorm(x+Sublayer(x))$, where $Sublayer(x)$ is the function implemented by the sub-layer itself. We apply dropout to the output of each sub-layer, before it is added to the sub-layer input and normalized.
class SublayerConnection(nn.Module):
"""
A residual connection followed by a layer norm.
Note for code simplicity the norm is first as opposed to last.
"""
def __init__(self, size, dropout):
super(SublayerConnection, self).__init__()
self.norm = LayerNorm(size)
self.dropout = nn.Dropout(dropout)
def forward(self, x, sublayer):
"Apply residual connection to any sublayer with the same size."
return x + self.dropout(sublayer(self.norm(x)))
逐位置前馈网络(Position-wise Feed-Forward Networks）
逐位置前馈网络(Position-wise Feed-Forward Networks)，为两层线性映射和他们之间的一个ReLU激活函数 。Another way of describing this is as two convolutions with kernel size 1。The dimensionality of input and output is $d_{model}=512$，and the inner-layer has dimensionality $d_{ff}=2048$.
class PositionwiseFeedForward(nn.Module):
def __init__(self, d_model, d_ff, dropout=0.1):
super(PositionwiseFeedForward, self).__init__()
self.w_1 = nn.Linear(d_model, d_ff) # 线性变换
self.w_2 = nn.Linear(d_ff, d_model) # 线性变换
self.dropout = nn.Dropout(dropout)
def forward(self, x):
x = self.w_1(x)
x = F.relu(x)
x = self.dropout(x)
x = self.w_2(x)
return x
编码器整体架构
基于上面的构建，我们来实现编码器层：
class EncoderLayer(nn.Module):
def __init__(self, size, self_attn, feed_forward, dropout):
super(EncoderLayer, self).__init__()
self.self_attn = self_attn # 多头注意力
self.feed_forward = feed_forward # 前馈网络
# SublayerConnection作用连接multi和ffn
self.sublayer = clones(SublayerConnection(size, dropout), 2)
# d_model
self.size = size
def forward(self, x, mask):
# 将embedding层进行多头注意力
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, mask)) # 先是多头注意力
# attn的结果直接作为下一层输入
return self.sublayer[1](x, self.feed_forward) # 然后是前馈网络
而编码器就是$N$个编码器层的叠加：
class Encoder(nn.Module):
def __init__(self, layer, N):
"""
layer = EncoderLayer
"""
super(Encoder, self).__init__()
# 复制N个编码器基本单元
self.layers = clones(layer, N)
# 层归一化
self.norm = LayerNorm(layer.size)
def forward(self, x, mask):
"""
循环编码器基本单元N次
"""
for layer in self.layers:
x = layer(x, mask) # 叠加N次
return self.norm(x) # 最后经过层归一化
解码器
解码器同样由N层解码器基本单元堆叠而成，与编码器基本单元不同的是：解码器在编码器基本单元的多头自注意力机制及前馈网络之间插入一个上下文注意力(context-attention机制Multi-Head Attention)层，用解码器基本单元的自注意力机制输出作为q查询编码器的输出，以便解码时，解码器获得编码器的所有输出，即上下文注意力机制的K和V来自编码器的输出，Q来自解码器前一时刻的输出。
解码器解码过程：给定编码器输出(编码器输入语句所有单词的词向量)和解码器前一时刻输出（单词），预测当前时刻单词的概率分布。
注意：
- 训练过程中，编、解码器均可以并行计算(训练语料中已知前一时刻单词)；
- 推理过程中，编码器可以并行计算，解码器需要像RNN一样依次预测输出单词。
解码器层的实现：
class DecoderLayer(nn.Module):
def __init__(self, size, self_attn, src_attn, feed_forward, dropout):
super(DecoderLayer, self).__init__()
self.size = size
# masked自注意力机制
self.self_attn = self_attn
# 上下文注意力机制
self.src_attn = src_attn
# 前馈网络
self.feed_forward = feed_forward
self.sublayer = clones(SublayerConnection(size, dropout), 3) # 解码器有三个子层
def forward(self, x, memory, src_mask, tgt_mask):
# memory为编码器输出隐藏表示
m = memory
# 自注意力机制，q、k、v均来自解码器隐表示 (子层一)
x = self.sublayer[0](x, lambda x: self.self_attn(x, x, x, tgt_mask))
# 上下文注意力机制：q为来自解码器隐表示，而k、v为编码器隐表示 (子层二)
x = self.sublayer[1](x, lambda x: self.src_attn(x, m, m, src_mask))
# 接下来是前馈网络(子层三)
return self.sublayer[2](x, self.feed_forward)
解码器的实现，主要是克隆了$N$个解码器层，最后的输出经过层归一化：
class Decoder(nn.Module):
def __init__(self, layer, N):
super(Decoder, self).__init__()
self.layers = clones(layer, N)
self.norm = LayerNorm(layer.size)
def forward(self, x, memory, src_mask, tgt_mask):
"""
循环解码器基本单元N次
"""
for layer in self.layers:
x = layer(x, memory, src_mask, tgt_mask)
return self.norm(x)
解码器最上面的部分作为生成器(generator)，由线性层+Softmax层组成：
class Generator(nn.Module):
"""
解码器输出经线性变换和softmax函数映射为下一时刻预测单词的概率分布
"""
def __init__(self, d_model, vocab):
super(Generator, self).__init__()
# decode后的结果，先进入一个全连接层变为词典大小的向量
self.proj = nn.Linear(d_model, vocab)
def forward(self, x):
# 然后再进行log_softmax操作(在softmax结果上再做多一次log运算)
return F.log_softmax(self.proj(x), dim=-1)
解码器注意力掩码:
解码器注意力掩码相对于编码器略微复杂，不仅需要将填充部分屏蔽掉,还需要对当前及后续序列进行屏蔽（subsequent_mask)，防止作弊。即解码器在预测当前时刻单词时，不能知道当前及后续单词内容，因此注意力掩码需要将当前时刻之后的注意力分数全部置为$-\infty$，然后再计算softmax，防止发生数据泄露。
def subsequent_mask(size):
"Mask out subsequent positions."
attn_shape = (1, size, size)
# 主对角线上移一位，主对角线下的元素全为0
subsequent_mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(
torch.uint8
)
return subsequent_mask == 0
subsequent_mask的矩阵形式为一个下三角矩阵，在主对角线右上位置全部为False。
注意力mask下面显示了每个tgt(目标)单词(行)允许查看的位置(列)。在训练过程中，当前单词后面的单词会被屏蔽。
plt.figure(figsize=(5,5)) plt.imshow(subsequent_mask(20)[0])
比如第0行，只能看到1列，第1行只能看到2列。黄色的区域代表可以看到的列。
Transformer模型整体架构
大部分有竞争力的神经网络序列转导模型都有一个编码器-解码器(Encoder-Decoder)结构。编码器映射一个用符号表示的输入序列$(x_1,\cdots,x_n)$到一个连续的序列表示$z=(z_1,\cdots, z_n)$。给定$z$，解码器生成符号的一个输出序列$ (y_1,\cdots,y_m)$，一次生成一个元素。在每个时间步，模型是自回归(auto-regressive)的，在生成下个输出时消耗上一次生成的符号作为附加的输入。
实际上Transformer就是一编码器-解码器架构，
class Transformer(nn.Module):
def __init__(self, encoder, decoder, src_embed, tgt_embed, generator):
super(Transformer, self).__init__()
self.encoder = encoder
self.decoder = decoder
self.src_embed = src_embed
self.tgt_embed = tgt_embed
self.generator = generator
def encode(self, src, src_mask):
return self.encoder(self.src_embed(src), src_mask)
def decode(self, memory, src_mask, tgt, tgt_mask):
return self.decoder(self.tgt_embed(tgt), memory, src_mask, tgt_mask)
def forward(self, src, tgt, src_mask, tgt_mask):
# encoder的结果作为decoder的memory参数传入，进行decode
return self.decode(self.encode(src, src_mask), src_mask, tgt, tgt_mask)
然后我们实现构建Transformer模型的函数：
def make_model(src_vocab, tgt_vocab, N=6, d_model=512, d_ff=2048, h = 8, dropout=0.1):
"Helper: Construct a model from hyperparameters."
c = copy.deepcopy
# 实例化Attention对象
attn = MultiHeadedAttention(h, d_model).to(DEVICE)
# 实例化FeedForward对象
ff = PositionwiseFeedForward(d_model, d_ff, dropout).to(DEVICE)
# 实例化PositionalEncoding对象
position = PositionalEncoding(d_model, dropout).to(DEVICE)
# 实例化Transformer模型对象
model = Transformer(
Encoder(EncoderLayer(d_model, c(attn), c(ff), dropout).to(DEVICE), N).to(DEVICE),
Decoder(DecoderLayer(d_model, c(attn), c(attn), c(ff), dropout).to(DEVICE), N).to(DEVICE),
nn.Sequential(Embeddings(d_model, src_vocab).to(DEVICE), c(position)),
nn.Sequential(Embeddings(d_model, tgt_vocab).to(DEVICE), c(position)),
Generator(d_model, tgt_vocab)).to(DEVICE)
# This was important from their code.
# Initialize parameters with Glorot / fan_avg.
for p in model.parameters():
if p.dim() > 1:
# 这里初始化采用的是nn.init.xavier_uniform
nn.init.xavier_uniform_(p)
return model.to(DEVICE)
模型训练及预测
标签平滑
训练过程中，采用KL散度损失实现标签平滑($\epsilon_{ls} = 0.1$)策略，提高模型鲁棒性、准确性和BLEU分数。
class LabelSmoothing(nn.Module):
"""
标签平滑
"""
def __init__(self, size, padding_idx, smoothing=0.0):
super(LabelSmoothing, self).__init__()
self.criterion = nn.KLDivLoss(reduction='sum')
self.padding_idx = padding_idx
self.confidence = 1.0 - smoothing
self.smoothing = smoothing
self.size = size
self.true_dist = None
def forward(self, x, target):
# x:[batch_size*seq_len,tgt_vocab]
# target:[batch_size*seq_len]
assert x.size(1) == self.size
# true_dist:[batch_size*seq_len,tgt_vocab]
true_dist = x.data.clone()
# 除了真正的类别位置和PAD位置，其它的位置smoothing平分
true_dist.fill_(self.smoothing / (self.size - 2))
true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
# PAD位置处的概率值设置为0，绝对不可能
true_dist[:, self.padding_idx] = 0
# mask:[z,1],当前位置下一个位置实际上是PAD的
mask = torch.nonzero(target.data == self.padding_idx)
if mask.dim() > 0:
# 如果当前位置的下一个位置实际上是PAD，则将这个位置的
# 的预测概率全部的设置为0
true_dist.index_fill_(0, mask.squeeze(), 0.0)
self.true_dist = true_dist
return self.criterion(x, true_dist.clone().detach())
计算损失
class SimpleLossCompute:
"""
简单的计算损失和进行参数反向传播更新训练的函数
"""
def __init__(self, generator, criterion, opt=None):
self.generator = generator
self.criterion = criterion # 损失函数，标签平滑后的版本
self.opt = opt # 优化器
def __call__(self, x, y, norm):
x = self.generator(x)
loss = self.criterion(x.contiguous().view(-1, x.size(-1)),
y.contiguous().view(-1)) / norm
loss.backward()# 进行反向传播操作，计算梯度
if self.opt is not None:
self.opt.step()
self.opt.optimizer.zero_grad()
return loss.data.item() * norm.float()
优化器
Adam优化器，$\beta_1=0.9、\beta_2=0.98$和$\epsilon = 10^{−9}$，并使用warmup策略调整学习率：
使用固定步数warmup_steps先使学习率的线性增长热身($step\_{num} \cdot warmup\_{steps}^{-1.5} = \frac {step\_{num}} {warmup\_{steps}} \cdot warmup\_{steps}^{-0.5}$），而后随着step_num的增加以step_num的反平方根成比例逐渐减小学习率。
主要调节是在rate 这个函数中，其中
- model_size 即为$d_{model}$
- warmup 即为warmup_steps
- factor可以理解为初始的学习率
class NoamOpt:
"Optim wrapper that implements rate."
def __init__(self, model_size, factor, warmup, optimizer):
self.optimizer = optimizer # 实际的优化器
self._step = 0 # 当前更新步数
self.warmup = warmup
self.factor = factor # 可以理解为初始的学习率
self.model_size = model_size
self._rate = 0
def step(self):
"Update parameters and rate"
self._step += 1
rate = self.rate()
for p in self.optimizer.param_groups:
p['lr'] = rate
self._rate = rate
self.optimizer.step()
def rate(self, step = None):#对学习率进行更新，返回新得到的学习率
"Implement `lrate` above"
if step is None:
step = self._step
return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))
def get_std_opt(model):
return NoamOpt(model.src_embed[0].d_model, 2, 4000,
torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))
训练
接下来，我们创建一个通用的训练和评分功能来跟踪损失。 我们传入一个上面定义的损失计算函数，它也处理参数更新，
def run_epoch(data, model, loss_compute, epoch):
start = time.time()
total_tokens = 0.
total_loss = 0.#总损失(sum形式)
tokens = 0.
for i , batch in enumerate(data):
out = model(batch.src, batch.tgt, batch.src_mask, batch.tgt_mask)
#这里得到的损失是sum形式的
loss = loss_compute(out, batch.tgt_y, batch.ntokens)
total_loss += loss
total_tokens += batch.ntokens
tokens += batch.ntokens
if i % 50 == 1:
elapsed = time.time() - start
print("Epoch %d Batch: %d Loss: %f Tokens per Sec: %fs" % (epoch, i + 1, loss / batch.ntokens, (tokens.float() / elapsed)))
start = time.time()
tokens = 0
return total_loss / total_tokens
然后就可以进行训练了
# 数据预处理
data = PrepareData(TRAIN_FILE, DEV_FILE)
src_vocab = len(data.en_word_dict)
tgt_vocab = len(data.cn_word_dict)
print("src_vocab %d" % src_vocab)
print("tgt_vocab %d" % tgt_vocab)
# 初始化模型
model = make_model(
src_vocab,
tgt_vocab,
LAYERS,
D_MODEL,
D_FF,
H_NUM,
DROPOUT
)
# 训练
print(">>>>>>> start train")
train_start = time.time()
criterion = LabelSmoothing(tgt_vocab, padding_idx = 0, smoothing= 0.0)
optimizer = NoamOpt(D_MODEL, 1, 2000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9,0.98), eps=1e-9))
train(data, model, criterion, optimizer)
print(f"<<<<<<< finished train, cost {time.time()-train_start:.4f} seconds")
预测
训练好了之后，我们用模型进行预测来看一下效果：
# 预测
# 加载模型
model.load_state_dict(torch.load(SAVE_FILE))
# 开始预测
print(">>>>>>> start evaluate")
evaluate_start = time.time()
evaluate(data, model)
print(f"<<<<<<< finished evaluate, cost {time.time()-evaluate_start:.4f} seconds")
模型预测用到的主要的函数为：
def greedy_decode(model, src, src_mask, max_len, start_symbol):
memory = model.encode(src, src_mask)
# start_symbol('BOS')的id
# ys:[1,1]
ys = torch.zeros(1, 1).fill_(start_symbol).type_as(src.data)
for i in range(max_len - 1):
out = model.decode(
memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
)
prob = model.generator(out[:, -1])
# next_word是一个列表
_, next_word = torch.max(prob, dim=1)
next_word = next_word.data[0]
ys = torch.cat(
[ys, torch.zeros(1, 1).type_as(src.data).fill_(next_word)], dim=1
)
return ys
def evaluate(data, model):
"""
在data上用训练好的模型进行预测，打印模型翻译结果
"""
# 梯度清零
with torch.no_grad():
# 在data的英文数据长度上遍历下标
for i in range(len(data.dev_en)):
# 打印待翻译的英文语句
en_sent = " ".join([data.en_index_dict[w] for w in data.dev_en[i]])
print("\n" + en_sent)
# 打印对应的中文语句答案
cn_sent = " ".join([data.cn_index_dict[w] for w in data.dev_cn[i]])
print("".join(cn_sent))
# 将当前以单词id表示的英文语句数据转为tensor，并放如DEVICE中
src = torch.from_numpy(np.array(data.dev_en[i])).long().to(DEVICE)
# 增加一维
src = src.unsqueeze(0)# src:[1,seq_len]
# 设置attention mask
src_mask = (src != 0).unsqueeze(-2)# src_mask:[1,1,seq_len]
# 用训练好的模型进行decode预测
out = greedy_decode(model, src, src_mask, max_len=MAX_LENGTH, start_symbol=data.cn_word_dict["BOS"])
# 初始化一个用于存放模型翻译结果语句单词的列表
translation = []
# 遍历翻译输出字符的下标（注意：开始符"BOS"的索引0不遍历）
for j in range(1, out.size(1)):
# 获取当前下标的输出字符
sym = data.cn_index_dict[out[0, j].item()]
# 如果输出字符不为'EOS'终止符，则添加到当前语句的翻译结果列表
if sym != 'EOS':
translation.append(sym)
# 否则终止遍历
else:
break
# 打印模型翻译输出的中文语句结果
print("translation: %s" % " ".join(translation))
预测效果如下：
BOS look around . EOS BOS 四 处 看 看 。 EOS translation: 看 ！ BOS hurry up . EOS BOS 赶 快 ! EOS translation: 快 点 ！ BOS keep trying . EOS BOS 继 续 努 力 。 EOS translation: 继 续 努 力 。 BOS take it . EOS BOS 拿 走 吧 。 EOS translation: 去 拿 吧 。 BOS birds fly . EOS BOS 鸟 类 飞 行 。 EOS translation: 鸟 类 飞 行 。 BOS hurry up . EOS BOS 快 点 ！ EOS translation: 快 点 ！ BOS look there . EOS BOS 看 那 里 。 EOS translation: 看 那 里 看 见 。 BOS how annoying ! EOS BOS 真 烦 人 。 EOS translation: 真 烦 人 。 BOS get serious . EOS BOS 认 真 点 。 EOS translation: 认 真 点 。 BOS once again . EOS BOS 再 一 次 。 EOS translation: 再 次 一 次 。 BOS stay sharp . EOS BOS 保 持 警 惕 。 EOS translation: 保 持 警 惕 。 BOS i won ! EOS BOS 我 赢 了 。 EOS translation: 我 赢 得 了 。 BOS get away ! EOS BOS 滚 ！ EOS translation: 走 开 ！ BOS i resign . EOS BOS 我 放 弃 。 EOS translation: 我 放 弃 。 BOS how strange ! EOS BOS 真 奇 怪 。 EOS translation: 真 奇 怪 。 ...
