Transformer训练机制
前言
- 关于Transformer原理与论文的介绍：详细了解Transformer：Attention Is All You Need
- PyTorch中实现Transformer模型
前面介绍了，Transformer 模型结构的实现，这里介绍下论文中提到的训练策略与设置。
设置文件名为training.py
Optimizer 优化器
文中选择 Adam 优化器，\(\beta_1=0.9, \beta_2=0.98, \epsilon=10^{-9}\)。学习率的变化函数定义为：
\[lrate=d_{\mathrm{model}}^{-0.5}\cdot\min(step\_num^{-0.5},step\_num\cdot warmup\_steps^{-1.5})
\]
在步数小于\(warmup\_steps\)时，学习率线性增长，之后学习率随着步数平方根倒数减少。
def rate(step, model_size, factor, warmup):
if step == 0:
step = 1
return factor * (
model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))
)
测试：学习率（y）随着steps（x）变化如图：
点击查看代码
import altair as alt
import pandas as pd
from training import rate
import torch
from torch.optim import Adam
from torch.optim.lr_scheduler import LambdaLR
def example_learning_schedule():
opts = [
[512, 1, 4000], [512, 1, 8000], [256, 1, 4000]
]
dummy_model = torch.nn.Linear(1, 1)
learning_rates = []
for idx, example in enumerate(opts): # 分别对三种样例进行测试
optimizer = Adam(dummy_model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9)
lr_scheduler = LambdaLR(optimizer=optimizer, lr_lambda=lambda step: rate(step, *example))
tmp = []
for step in range(20000):
tmp.append(optimizer.param_groups[0]["lr"])
optimizer.step()
lr_scheduler.step()
learning_rates.append(tmp)
learning_rates = torch.tensor(learning_rates) # 不转为张量报错： TypeError: list indices must be integers or slices, not tuple
alt.data_transformers.disable_max_rows() # 禁用最大行限制
opts_data = pd.concat([
pd.DataFrame({
"Learning Rate": learning_rates[warmup_idx, :],
"model_size - warmup": ["512 - 4000", "512 - 8000", "256 - 4000"][warmup_idx],
"step": range(20000)
})
for warmup_idx in [0, 1, 2]
])
return (
alt.Chart(opts_data)
.mark_line()
.properties(width=600)
.encode(x="step", y="Learning Rate", color="model_size - warmup:N")
.interactive()
)
example_learning_schedule()
Regularization 正则化
Dropout
将 dropout 引用于每个子层的输出：输入 x 进行归一化 -> dropout -> 残差连接。还在 Encoder 和 Decoder 的嵌入层与位置编码层的输出结果上使用了 dropout。基础模型上 dropout 的概率为\(P_{drop}=0.1\)。
代码中通过使用 PyTorch 下的nn.Dropout()实现 dropout。
nn.Dropout()初始化参数p表示训练时，以概率 p 将输入张量的一些元素归零，对于没有归零的元素将乘以\(\frac{1}{1-p}\)。- 输入为任意形状的张量，输出为与输入张量形状相同并经过处理的张量。[Source]
Label Smoothing 标签平滑
Label Smoothing 可以帮助模型提高泛化，减轻模型“过度自信”，具体来说，对于每个训练样本，标签平滑会将正确的标签设置为一个稍微小于1的值，同时将错误的标签设置为一个较小的非零值，为模型提供更丰富的信息。论文中，设置平滑度\(\epsilon_{ls}=0.1\)。
class LabelSmoothing(nn.Module):
# size 词典大小
# padding_idx 填充索引
# smoothing 平滑值
def __init__(self, size, padding_idx, smoothing=0.0):
super(LabelSmoothing, self).__init__()
self.criterion = nn.KLDivLoss(reduction="sum") # KL散度损失， ‘sum’ 表示对所有样本KL散度求和
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):
assert x.size(1) == self.size
true_dist = x.data.clone()
# 对 true_dist 进行填充，填充值为 smoothing / (size - 2)
true_dist.fill_(self.smoothing / (self.size - 2))
# 将 target 中的每个元素作为索引，将 true_dist 中的对应元素置为 confidence
true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
# 将 target 中的 padding_idx 对应的元素置为 0
true_dist[:, self.padding_idx] = 0
# 找到 target 中为 padding_idx 的元素的索引
mask = torch.nonzero(target.data == self.padding_idx)
# 如果 mask 不为空，将 true_dist 中对应的元素置为 0（掩码位置处取0）
if mask.dim() > 0:
true_dist.index_fill_(0, mask.squeeze(), 0.0)
self.true_dist = true_dist
# 计算模型输出 x 与平滑处理后的 true_dist 之间的 KL 散度损失
return self.criterion(x, true_dist.clone().detach())
对于 PyTorch 下的张量，使用clone()断开内存的联系，detach() 断开计算图的连接，使张量不参与梯度计算（顺序可互换）。
测试：行表示单词，列表示单词在对应标签的概率。
点击查看代码
import altair as alt
import pandas as pd
from training import LabelSmoothing
import torch
def example_label_smoothing():
crit = LabelSmoothing(5, 0, 0.4)
predict = torch.FloatTensor([
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0],
[0, 0.2, 0.7, 0.1, 0]
])
crit(x=predict, target=torch.LongTensor([2, 1, 0, 3, 3]))
LS_data = pd.concat([
pd.DataFrame({
"target distribution": crit.true_dist[x, y].flatten(),
"columns": y,
"rows": x
})
for x in range(5)
for y in range(5)
])
return (
alt.Chart(LS_data)
.mark_rect()
.properties(height=200, width=200)
.encode(
alt.X("columns:O"),
alt.Y("rows:O"),
alt.Color("target distribution:Q", scale=alt.Scale(scheme="viridis")),
)
)
example_label_smoothing()
mask 掩码
Decoder 的自注意力层为了防止模型关注位置之后的信息增加了掩码：Outputs 的嵌入词右偏移一个位置，并让位置 i 的预测只能依赖于小于 i 的位置的已知输出。
掩码的添加放在训练 Batch 的创建中，src表示输入序列，tgt表示输出序列（由于测试时输出序列未知，因此tgt可设置为 None），pad表示填充的索引（填充的值由自己确定，一般取2）。
# return Shape: 1 * size * size
def subsequent_mask(size):
attn_shape = (1, size, size)
mask = torch.triu(torch.ones(attn_shape), diagonal=1).type(torch.uint8) # 上三角矩阵
return mask == 0
class Batch:
def __init__(self, src, tgt=None, pad=2):
self.src = src
self.src_mask = (src != pad).unsqueeze(-2)
if tgt is not None:
self.tgt = tgt[:, :-1]
self.tgt_y = tgt[:, 1:]
self.tgt_mask = self.make_std_mask(self.tgt, pad)
self.ntokens = (self.tgt_y != pad).data.sum()
# return Shape: batch_size * tgt.size(-2) * tgt.size(-1)
@staticmethod # 在 pad 处添加掩码
def make_std_mask(tgt, pad):
tgt_mask = (tgt != pad).unsqueeze(-2)
tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)
return tgt_mask
测试：返回的mask矩阵形状：序列长度 * 序列长度。
点击查看代码
import altair as alt
import pandas as pd
from training import subsequent_mask
def example_mask():
size = 20
LS_data = pd.concat( # Shape: (400, 3)
[
pd.DataFrame(
{
"Subsequent Mask": subsequent_mask(size)[0][x, y].flatten(),
"Window": y,
"Masking": x,
}
)
for y in range(20)
for x in range(20)
]
)
return (
alt.Chart(LS_data)
.mark_rect()
.properties(height=250, width=250)
.encode(
alt.X("Window:O"),
alt.Y("Masking:O"),
alt.Color("Subsequent Mask:Q", scale=alt.Scale(scheme="viridis")),
)
.interactive()
)
example_mask()
Loss 损失函数
损失函数的计算已经包括在 Label Smoothing 中，这里只封装了一个损失模块：
class SimpleLossCompute:
def __init__(self, criterion):
self.criterion = criterion
def __call__(self, x, y, norm):
sloss = ( # sloss: scaled loss
self.criterion(
x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)
) / norm
)
# .data 表示没有梯度信息的新的张量，其数据与原始张量共享底层数据，
# 但对这个新的张量进行操作不会影响梯度传播回原始张量
return sloss.data * norm, sloss
测试：损失变化。
点击查看代码
from training import LabelSmoothing
import torch.nn.functional as F
import torch
import pandas as pd
import altair as alt
def loss(x, crit):
d = x + 3 * 1
predict = torch.FloatTensor([[0, x/d, 1/d, 1/d, 1/d]])
return crit(F.log_softmax(predict), torch.LongTensor([1])).data
def penalization_visualization():
crit = LabelSmoothing(5, 0, 0.1)
loss_data = pd.DataFrame({
"Loss": [loss(x, crit) for x in range(1, 100)],
"Steps": list(range(99))
}).astype("float")
# print(loss_data["Loss"])
return (
alt.Chart(loss_data)
.mark_line()
.properties(width=350)
.encode(x="Steps", y="Loss")
.interactive()
)
penalization_visualization()
Decode 解码
对于最后模型生成器的输出张量，使用贪心解码：在每个时间步选择当前概率最高的单词或符号作为输出，然后将其作为下一个时间步的输入，依次生成整个序列，直到生成结束符号或达到预设的最大长度。
def greedy_decode(model, src, src_mask, max_len, start_symbol):
memory = model.encode(src, src_mask)
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 = 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
相关环境
altair 5.2.0
pandas 2.1.4
torch 2.1.2
