Transformer训练细节

完整代码:The Annotated Transformer

前言

  1. 关于Transformer原理与论文的介绍:详细了解Transformer:Attention Is All You Need
  2. PyTorch中实现Transformer模型

前面介绍了,Transformer 模型结构的实现,这里介绍下论文中提到的训练策略与设置。
设置文件名为training.py

Optimizer 优化器

文中选择 Adam 优化器,β1=0.9,β2=0.98,ϵ=109。学习率的变化函数定义为:

lrate=dmodel0.5min(step_num0.5,step_numwarmup_steps1.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)变化如图:

相关测试代码(在 Jupyter Notebook 中使用 altair 库绘制):
点击查看代码
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 的概率为Pdrop=0.1

代码中通过使用 PyTorch 下的nn.Dropout()实现 dropout。

  • nn.Dropout()初始化参数p表示训练时,以概率 p 将输入张量的一些元素归零,对于没有归零的元素将乘以11p
  • 输入为任意形状的张量,输出为与输入张量形状相同并经过处理的张量。[Source]

Label Smoothing 标签平滑

Label Smoothing 可以帮助模型提高泛化,减轻模型“过度自信”,具体来说,对于每个训练样本,标签平滑会将正确的标签设置为一个稍微小于1的值,同时将错误的标签设置为一个较小的非零值,为模型提供更丰富的信息。论文中,设置平滑度ϵ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() 断开计算图的连接,使张量不参与梯度计算(顺序可互换)。

测试:行表示单词,列表示单词在对应标签的概率。

相关测试代码(在 Jupyter Notebook 中使用 altair 库绘制):
点击查看代码
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表示<blank>的索引,<blank>是特殊token之一表示空格用于填充序列,本文中特殊token列表为specials=["<s>", "</s>", "<blank>", "<unk>"],因此pad=2表示空格。src != pad表示将输入序列中不为 pad 的位置设置为 True,pad位置设置为 False,这样就可以在后续计算中忽略 pad 位置的信息。

# 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 # 上三角部分设置为 False,下三角部分设置为 True
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

注意:transformer架构图中,共有3处计算 multi-head attention,其中1处为 masked multi-head attention。实际上,每处 multi-head attention 计算都使用了mask,图中“Multi-Head Attention”表示的是仅仅对 padding的部分mask:

self.src_mask = (src != pad).unsqueeze(-2)

而“Masked Multi-Head Attention”表示的是对 padding和未来的部分mask:

tgt_mask = (tgt != pad).unsqueeze(-2)
tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(tgt_mask.data)

这里,encoder和decoder中的“Multi-Head Attention”都使用了src_mask,而decoder中的“Masked Multi-Head Attention”使用了tgt_mask

测试:返回的mask矩阵形状:序列长度 * 序列长度。

相关测试代码(在 Jupyter Notebook 中使用 altair 库绘制):
点击查看代码
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

测试:损失变化。

相关测试代码(在 Jupyter Notebook 中使用 altair 库绘制):
点击查看代码
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

参考文献

  1. The Annotated Transformer
posted @   October-  阅读(310)  评论(0编辑  收藏  举报
相关博文:
阅读排行:
· 清华大学推出第四讲使用 DeepSeek + DeepResearch 让科研像聊天一样简单!
· 推荐几款开源且免费的 .NET MAUI 组件库
· 实操Deepseek接入个人知识库
· 易语言 —— 开山篇
· Trae初体验
历史上的今天:
2023-02-07 爬取天气信息
点击右上角即可分享
微信分享提示