Transformer训练细节
完整代码:The Annotated Transformer
前言
- 关于Transformer原理与论文的介绍:详细了解Transformer:Attention Is All You Need
- PyTorch中实现Transformer模型
前面介绍了,Transformer 模型结构的实现,这里介绍下论文中提到的训练策略与设置。
设置文件名为training.py
Optimizer 优化器
文中选择 Adam 优化器,
在步数小于
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 的概率为
代码中通过使用 PyTorch 下的nn.Dropout()
实现 dropout。
nn.Dropout()
初始化参数p
表示训练时,以概率 p 将输入张量的一些元素归零,对于没有归零的元素将乘以 。- 输入为任意形状的张量,输出为与输入张量形状相同并经过处理的张量。[Source]
Label Smoothing 标签平滑
Label Smoothing 可以帮助模型提高泛化,减轻模型“过度自信”,具体来说,对于每个训练样本,标签平滑会将正确的标签设置为一个稍微小于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表示<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矩阵形状:序列长度 * 序列长度。
点击查看代码
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
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 清华大学推出第四讲使用 DeepSeek + DeepResearch 让科研像聊天一样简单!
· 推荐几款开源且免费的 .NET MAUI 组件库
· 实操Deepseek接入个人知识库
· 易语言 —— 开山篇
· Trae初体验
2023-02-07 爬取天气信息