万字逐行解析与实现Transformer，并进行德译英实战（二）
文章目录
- Part 2: 模型训练
-
- Batches and Masking
- Training Loop
- 训练数据和Batch
- Optimizer
- 正则化
-
- 标签平滑
- 第一个例子
-
- 造数据
- 损失计算
- 使用贪心算法解码(Greedy Decoding)
本文由于长度限制，共分为三篇：
- 万字逐行解析与实现Transformer，并进行德译英实战（一）
- 万字逐行解析与实现Transformer，并进行德译英实战（二）
- 万字逐行解析与实现Transformer，并进行德译英实战（三）
你也可以在该项目找到本文的源码。
Part 2: 模型训练
本章描述模型的训练机制。
在正式进入训练模型之前，需要先定义一些工具类和方法。
Batches and Masking
class Batch:
"""
定义一个Batch，来存放一个batch的src，tgt，src_mask等对象。
方便后续的取用
"""
def __init__(self, src, tgt=None, pad=2): # 2 = ”、“”和“”是要算作句子成分的）
这里unsqueeze一下是因为后续是要对Attention中的scores进行mask，而scores的len(shape)=3,
为了与scores保持一致，所以unsqueeze(-2)一下。具体可参考attention函数中的注释。
"""
self.src_mask = (src != pad).unsqueeze(-2)
if tgt is not None:
"""
每个句子都去掉最后一个词。例如tgt的Shape为(16, 30)，
即batch_size为16，每个句子30个单词，执行该代码后为
(16, 29)。这么做的原因是：tgt存储的是Decoder的输入，
而Decoder的输入是不可能包含最后一个词的。例如，我们要
预测` 我 爱 你 `，第一次我们会给Decoder
传``，第二次传` 我`，而最后一次会传
` 我 爱 你`。所以tgt不会出现最后一个词，所以要去掉。
"""
self.tgt = tgt[:, :-1]
"""
与上面差不多，去掉句子的第一个词，也就是“”
tgt_y 存储的是希望预测的结果，所以不需要''
例如，传入encoder的是" I love you "，
初始传入decoder为""，则我们想最终能“一个个”
的预测出“我 爱 你 ”，所以要把去掉，因为
其不是我们想要预测的token。
"""
self.tgt_y = tgt[:, 1:]
"""
构造tgt_mask：tgt_mask与src_mask略有不同，除了需要盖住pad部分，还需要
将对角线右上的也都盖住，具体原因可参考：https://blog.csdn.net/zhaohongfei_358/article/details/125858248
例如：[[ 0, 5, 4, 6, 1, 2, 2 ]]，则tgt_mask则为：
[[[ True, False, False, False, False, False, False],
[ True, True, False, False, False, False, False],
[ True, True, True, False, False, False, False],
[ True, True, True, True, False, False, False],
[ True, True, True, True, True, False, False],
[ True, True, True, True, True, False, False],
[ True, True, True, True, True, False, False]]]
"""
self.tgt_mask = self.make_std_mask(self.tgt, pad)
"""
此Batch的tgt_y的总token数量。值为一个数字，例如 tensor(266) 表示tgt_y有266个有意义的token
也就是出去部分的词的数量。保存这个是为了用于loss的正则化，后面loss部分会详细说明。
注意这里是tgt_y，而tgt_y去掉了“”，所以token数是不包含“”的。
"""
self.ntokens = (self.tgt_y != pad).data.sum()
@staticmethod
def make_std_mask(tgt, pad=2):
"""
生成tgt_mask
"""
# 首先生成非句子成分部分的mask
# 例如 [[ 0, 5, 4, 6, 1, 2, 2 ]] 的mask为 [[[ True, True, True, True, True, False, False ]]]
tgt_mask = (tgt != pad).unsqueeze(-2)
"""
subsequent_mask用于获取阶梯式Mask。然后再和tgt_mask进行&操作。
例如：tgt_mask为[[[ True, True, False ]]]
subsequent_mask结果为：
[[[ True, False, False ],
[ True, True, False ],
[ True, True, True ]]]
则tgt_mask & subsequent_mask结果为：
[[[ True, False, False ],
[ True, True, False ],
[ True, True, False ]]]
"""
tgt_mask = tgt_mask & subsequent_mask(tgt.size(-1)).type_as(
tgt_mask.data
)
return tgt_mask
Training Loop
接下来我们定义一个通用的训练函数和训练状态类：
class TrainState:
"""用于保存一些训练状态"""
# step的次数，但注意是一次loss.backward()算一次，或者说一个batch算一次
# 并不是一次optimizer.step()算一次。在后面的训练代码中，可能会累计多次loss
# 然后进行一次optimizer.step()
step: int = 0
# 参数更新的次数。这个才是optimizer.step()的次数
accum_step: int = 0
samples: int = 0 # 记录训练过的样本数量
tokens: int = 0 # 记录处理过的token数量（target的）
def run_epoch(
data_iter,
model,
loss_compute,
optimizer,
scheduler,
mode="train",
accum_iter=1,
train_state=TrainState(),
):
"""
进行一个epoch训练
data_iter: 可迭代对象，一次返回一个Batch对象
model: Transformer模型，EncoderDecoder类对象
loss_compute: SimpleLossCompute对象，用于计算损失
optimizer: Adam优化器。验证时，optimizer是DummyOptimizer
scheduler：LambdaLR对象，用于调整Adam的学习率，实现WarmUp
若对调整学习率不熟悉，可参考：https://blog.csdn.net/zhaohongfei_358/article/details/125759911
验证时，scheduler是DummyScheduler
accum_iter: 多少个batch更新一次参数，默认为1，也就是每个batch都对参数进行更新
train_state: TrainState对象，用于保存一些训练状态
"""
start = time.time()
# 记录target的总token数，每次打印日志后，进行清0
tokens = 0
# 记录tgt_y的总token数，用于对total_loss进行正则化
total_tokens = 0
total_loss = 0
n_accum = 0 # 本次epoch更新了多少次模型参数
for i, batch in enumerate(data_iter):
# 前向传递。等价于model(batch.src, batch.tgt, batch.src_mask, batch.tgt_mask)
# 但注意，这里的out是Decoder的输出，并不是Generator的输出，因为在EncoderDecoder
# 的forward中并没有使用generator。generator的调用放在了loss_compute中
out = model.forward(
batch.src, batch.tgt, batch.src_mask, batch.tgt_mask
)
"""
计算损失，传入的三个参数分别为：
1. out: EncoderDecoder的输出，该值并没有过最后的线性层，过线性层被集成在了计算损失中
2. tgt_y: 要被预测的所有token，例如src为` I love you `，则`tgt_y`则为
`我 爱 你 `
3. ntokens：这批batch中有效token的数量。用于对loss进行正则化。
返回两个loss，其中loss_node是正则化之后的，所以梯度下降时用这个。
而loss是未进行正则化的，用于统计total_loss。
"""
loss, loss_node = loss_compute(out, batch.tgt_y, batch.ntokens)
# loss_node = loss_node / accum_iter
if mode == "train" or mode == "train+log":
# 计算梯度
loss_node.backward()
# 记录step次数
train_state.step += 1
# 记录样本数量。batch.src.shape[0]获取的是Batch size
train_state.samples += batch.src.shape[0]
# 记录处理过的token数
train_state.tokens += batch.ntokens
# 如果达到了accum_iter次，就进行一次参数更新
if i % accum_iter == 0:
optimizer.step()
optimizer.zero_grad(set_to_none=True)
# 记录本次epoch的参数更新次数
n_accum += 1
# 记录模型的参数更新次数
train_state.accum_step += 1
# 更新学习率
scheduler.step()
# 累计loss
total_loss += loss
# 累计处理过的tokens
total_tokens += batch.ntokens
# 累计从上次打印日志开始处理过得tokens
tokens += batch.ntokens
# 每40个batch打印一次日志。
if i % 40 == 1 and (mode == "train" or mode == "train+log"):
# 打印一下当前的学习率
lr = optimizer.param_groups[0]["lr"]
# 记录这40个batch的消耗时间
elapsed = time.time() - start
# 打印日志
print(
(
"Epoch Step: %6d | Accumulation Step: %3d | Loss: %6.2f "
+ "| Tokens / Sec: %7.1f | Learning Rate: %6.1e"
)
# i: 本次epoch的第几个batch
# n_accum: 本次epoch更新了多少次模型参数
# loss / batch.ntokens: 对loss进行正则化，然后再打印loss，其实这里可以直接用loss_node
# tokens / elapsed: 每秒可以处理的token数
# lr: 学习率（learning rate），这里打印学习率的目的是看一下warmup下学习率的变化
% (i, n_accum, loss / batch.ntokens, tokens / elapsed, lr)
)
# 重置开始时间
start = time.time()
# 重置token数
tokens = 0
del loss
del loss_node
# 返回正则化之后的total_loss，返回训练状态
return total_loss / total_tokens, train_state
训练数据和Batch
我们在标准的WMT 2014 English-German数据集上进行了训练，该数据集包含450w个句子对儿。句子使用的是byte-pair编码（subword的一种方法，如果你对subword不了解，可以参考这篇文章），其中source和target使用的是同一个词典。对于英文到法文的翻译，我们使用了一个巨大的WMT 2014 English-French数据集，其包含3600W个句子，其能拆分出的词典大小为32000。
具有相似长度的句子对儿将会组成一个Batch。每个训练batch都会包含一组句子对儿，其包含大约25000个source tokens和25000个target tokens.
Optimizer
Transformer使用了Warmup的方式来调整学习率，公式如下：
l r a t e = d model − 0.5 ⋅ min ( s t e p _ n u m − 0.5 , s t e p _ n u m ⋅ w a r m u p _ s t e p s − 1.5 ) lrate = d_{\text{model}}^{-0.5} \cdot \min({step\_num}^{-0.5}, {step\_num} \cdot {warmup\_steps}^{-1.5}) lrate=dmodel−0.5⋅min(step_num−0.5,step_num⋅warmup_steps−1.5)
等价于：
l r a t e = 1 d m o d e l m i n ( 1 s t e p _ n u m , s t e p _ n u m ⋅ 1 w a r m u p _ s t e p s w a r m u p _ s t e p s ) lrate = \frac{1}{\sqrt{d_{model}}} min(\frac{1}{\sqrt{step\_num}}, step\_num \cdot \frac{1}{warmup\_steps \sqrt{warmup\_steps}}) lrate=dmodel1min(step_num1,step_num⋅warmup_stepswarmup_steps1)
其中：
- l r a t e lrate lrate: 学习率的调整参数。注意这个并不是学习率，假设你设置的学习率为0.8，则第 i i i步的学习率为： 0.8 ∗ l r a t e ( i ) 0.8 * lrate(i) 0.8∗lrate(i)。
- d m o d e l d_{model} dmodel：模型的维度，即词向量的维度
- s t e p _ n u m step\_num step_num: 步数。注意：执行一个backward（也可以说是一个batch）step加1，并不是optimizer.step()一次step加1.
- warmup_steps: warmup多少步。在Transformer中使用的是4000，也就是预热4000步。
# 学习率调整函数
def rate(step, model_size, factor, warmup):
# 避免分母为0
if step == 0:
step = 1
# 这里比原公式还多了一个factor，factor默认取1，相当于没有多。
return factor * (
model_size ** (-0.5) * min(step ** (-0.5), step * warmup ** (-1.5))
)
接下来是一个学习率调整的例子，如果你不熟悉自定义学习率，可以参考这篇文章。
def example_learning_schedule():
# 准备3个样例，三个参数分别为：d_model, factor, warmup_steps
opts = [
[512, 1, 4000], # example 1
[512, 1, 8000], # example 2
[256, 1, 4000], # example 3
]
# 随便定义一个没啥用的模型
dummy_model = torch.nn.Linear(1, 1)
# 记录这三个例子的学习率变化
learning_rates = []
# 分别运行上面的三个例子
for idx, example in enumerate(opts):
# 定义Adam优化器
optimizer = torch.optim.Adam(
dummy_model.parameters(), lr=1, betas=(0.9, 0.98), eps=1e-9
)
# 定义LambdaLR调整学习率。"*example"称为序列解包
lr_scheduler = LambdaLR(
optimizer=optimizer, lr_lambda=lambda step: rate(step, *example)
)
# 用于存储当前样例的学习率变化
tmp = []
# 进行20000次训练
for step in range(20000):
# 记录学习率
tmp.append(optimizer.param_groups[0]["lr"])
# 参数更新
optimizer.step()
# 更新学习率。更新学习率要放在optimizer.step()之后
lr_scheduler.step()
# 记录当前样例的学习率
learning_rates.append(tmp)
# 将学习率变化变成
learning_rates = torch.tensor(learning_rates)
# Enable altair to handle more than 5000 rows
alt.data_transformers.disable_max_rows()
"""
组织统计图需要的数据，最终效果为：
|Learning Rate | model_size:warmup | step
----------------------------------------------
0 | 1.74e-07 | 512:4000 | 0
1 | 1.74e-07 | 512:4000 | 1
2 | 3.49e-07 | 512:4000 | 2
....
"""
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()
正则化
标签平滑
在Transformer的训练中，使用Label Smoothing技术对标签做了平滑处理，这样可以减少模型overconfidence，从而减少overfitting。
公式为：
P i = { 1 if ( i = y ) 0 if ( i ≠ y ) ) ⇒ P i = { 1 − ϵ if ( i = y ) ϵ K − 1 if ( i ≠ y ) ) P_{i} = \begin{cases} 1 ~~~&\text{if } (i=y) \\ 0 ~~~&\text{if } (i\neq y)) \end{cases} ⇒ P_{i} = \begin{cases} 1-\epsilon ~~~&\text{if } (i=y) \\ \frac{\epsilon}{K-1} ~~~&\text{if } (i\neq y)) \end{cases} Pi={1 0 if (i=y)if (i=y))⇒Pi={1−ϵ K−1ϵ if (i=y)if (i=y))
其中 P i P_i Pi 是标签的第i维的值； K K K为类别总数； ϵ \epsilon ϵ为平滑因子，取值范围为 [ 0 , 1 ] [0,1] [0,1], 通常为0.1，
例如，假设我们的标签是2, 词典大小为5，则对应的one-hot向量则为： [ 0 , 0 , 0 , 1 , 0 ] [0, 0, 0, 1, 0] [0,0,0,1,0]，我们现在取平滑因子 ϵ = 0.2 \epsilon=0.2 ϵ=0.2，则平滑后的Label为：
[ 0.2 / 4 , 0.2 / 4 , 0.2 / 4 , 1 − 0.2 , 0.2 / 4 ] = [ 0.05 , 0.05 , 0.05 , 0.8 , 0.05 ] [0.2/4, 0.2/4, 0.2/4, 1-0.2, 0.2/4]=[0.05, 0.05, 0.05, 0.8, 0.05] [0.2/4,0.2/4,0.2/4,1−0.2,0.2/4]=[0.05,0.05,0.05,0.8,0.05]
之后计算loss时，使用平滑后的标签。直观上的理解就是：模型你别太自信，就算你预测对了也对你进行一点惩罚，防止你过度相信这个特定样本的特定结果。
不过在下面的LabelSmoothing的实现与上面说的略有不同，主要有两点：
- 这个LabelSmoothing类将损失计算也放了进来，也就是它除了负责标签平滑外，还负责计算损失。
- 由于词典中有一项是“填充”(
在本代码实现中，所有的
指都是
对于第二点，做个详细的说明。
假设我们的词典为{0: ，此时词典大小为2000，所以one-hot向量的维度也是2000。如果我们做预测，无论什么情况我们的标签都不应该是
class LabelSmoothing(nn.Module):
"""
实现标签平滑。
该类除了标签平滑外，还包含了损失的计算。
所以此类为 LabelSmoothing + LossFunction
"""
def __init__(self, size, padding_idx, smoothing=0.0):
"""
size: 目标词典的大小。
padding_idx: 空格('')在词典中对应的index，``等价于``
smoothing: 平滑因子，0表示不做平滑处理
"""
super(LabelSmoothing, self).__init__()
# 定义损失函数，这里使用的是KL Divergence Loss，也是一种多分类常用的损失函数
# KLDivLoss官方文档：https://pytorch.org/docs/stable/generated/torch.nn.KLDivLoss.html
self.criterion = nn.KLDivLoss(reduction="sum")
self.padding_idx = padding_idx
self.confidence = 1.0 - smoothing
self.smoothing = smoothing
self.size = size
# true distribution，平滑后的标签
self.true_dist = None
def forward(self, x, target):
"""
x: generator输出的概率分布。Shape为(batch_size, 词典大小)
target: 目标标签。Shape为(batch_size)
"""
# 确保generator的输出维度和词典大小一致，否则后面计算loss的时候就会出错
assert x.size(1) == self.size
# 这句其实是为了创建一个与x有相同Shape的tensor。
# 这里假设x的shape为(2, 6)，即batch_size为2，词典大小为6。
true_dist = x.data.clone()
"""
将true_dist全部填充为 self.smoothing / (self.size - 2)。
假设 smoothing=0.2，则为全部填充为 0.2 / 4= 0.05
此时true_dist则为：
[[0.05, 0.05, 0.05, 0.05, 0.05, 0.05],
[0.05, 0.05, 0.05, 0.05, 0.05, 0.05]]
"""
true_dist.fill_(self.smoothing / (self.size - 2))
"""
scatter_: 官方地址 https://pytorch.org/docs/stable/generated/torch.Tensor.scatter_.html
将true_dist的1维度上与target.data.unsqueeze(1)对应的值变为src。
假设此例中target.data.unsqueeze(1) 为[[2], [3]]，即2个数据的标签分别为2，3
则true_dist执行过scatter后变为:
[[0.05, 0.05, 0.8, 0.05, 0.05, 0.05],
[0.05, 0.05, 0.05, 0.8, 0.05, 0.05]]
"""
true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
# 将"空格"所在的index填充为0。
true_dist[:, self.padding_idx] = 0
# 找出target中，label为的标签。例如target为['i', 'love', 'you', '', '']
# 那么mask则为[[1], [3]]，表示第1个和第3个为空格。
# 但在实际应用场景中标签中不应该出现""
mask = torch.nonzero(target.data == self.padding_idx)
if mask.dim() > 0:
"""
将""所在的label整个设置为0。
例如，假设现在true_dist为：
[[0.05, 0.05, 0.8, 0.05, 0.05, 0.05],
[0.05, 0.05, 0.05, 0.8, 0.05, 0.05],
[0.8, 0.05, 0.05, 0.05, 0.05, 0.05]]
其中第3行为“”的label，则执行完这行代码后，则为：
[[0.05, 0.05, 0.8, 0.05, 0.05, 0.05],
[0.05, 0.05, 0.05, 0.8, 0.05, 0.05],
[0.00, 0.00, 0.00, 0.0, 0.00, 0.00]]
"""
true_dist.index_fill_(0, mask.squeeze(), 0.0)
# 保存平滑标签后的label
self.true_dist = true_dist
"""
使用平滑后的标签计算损失
由于对``部分进行了mask，所以在这部分是不会参与损失计算的
"""
return self.criterion(x, true_dist.clone().detach())
这里对来看一个Label smoothing的例子：
def example_label_smoothing():
# 定义LabelSmoothing, 词典大小为5， 对应index为0，平滑因子为0.4
crit = LabelSmoothing(5, 0, 0.4)
# 定义predict，即网络的预测分布。
# 这里的1e-9在源码中为0。由于KLDivLoss需要对predict求log，而0会导致结果为负无穷
# 所以我这里将0改成了1e-9。
predict = torch.FloatTensor(
[
[1e-9, 0.2, 0.7, 0.1, 1e-9],
[1e-9, 0.2, 0.7, 0.1, 1e-9],
[1e-9, 0.2, 0.7, 0.1, 1e-9],
[1e-9, 0.2, 0.7, 0.1, 1e-9],
[1e-9, 0.2, 0.7, 0.1, 1e-9],
]
)
loss = crit(x=predict.log(), target=torch.LongTensor([2, 1, 0, 3, 3]))
print("loss:", loss)
print("Before label smoothing:\n", torch.zeros(5, 6).scatter_(1, torch.LongTensor([2, 1, 0, 3, 3]).unsqueeze(1), 1))
print("After label smoothing:\n", crit.true_dist)
LS_data = pd.concat(
[
pd.DataFrame(
{
"target distribution": crit.true_dist[x, y].flatten(),
"columns": y,
"rows": x,
}
)
for y in range(5)
for x in range(5)
]
)
return (
alt.Chart(LS_data)
.mark_rect(color="Blue", opacity=1)
.properties(height=200, width=200)
.encode(
alt.X("columns:O", title=None),
alt.Y("rows:O", title=None),
alt.Color(
"target distribution:Q", scale=alt.Scale(scheme="viridis")
),
)
.interactive()
)
example_label_smoothing()
上图中越黄，表示这个位置的值越接近1, 反之则越接近0。没平滑前，应该是黄色的部分巨黄（=1），其他部位都是巨紫（=0），平滑之后，黄色部位的黄色素给其他紫的部分匀了一些（除了’
'位置）。
下面是另一个平滑的例子，从该例子可以看到当模型非常自信的时候就会给予其一个微小的惩罚：
def loss(x, crit):
# x是从0到100的一个不断增大的数。 d=x+3，比x大一点。
d = x + 3
"""
模拟模型的输出。
一开始x为1，模型输出为：[[0.0000, 0.2500, 0.2500, 0.2500, 0.2500]]
此时模型模型还不太会预测
当x到100时，模型输出为：[[0.0000, 0.9706, 0.0098, 0.0098, 0.0098]]
此时模型可以很自信的说结果就是 1
"""
predict = torch.FloatTensor([[0, x / d, 1 / d, 1 / d, 1 / d]])
# 计算模型损失。由于使用的是KLDivLoss，所以要对predict进行log操作
return crit(predict.log(), torch.LongTensor([1])).data
def penalization_visualization():
crit = LabelSmoothing(5, 0, 0.1)
loss_data = pd.DataFrame(
{
# x从1开始不断增大，模拟模型的表现越来越好
"Loss": [loss(x, crit) for x in range(1, 100)],
"Steps": list(range(99)),
}
).astype("float")
return (
alt.Chart(loss_data)
.mark_line()
.properties(width=350)
.encode(
x="Steps",
y="Loss",
)
.interactive()
)
penalization_visualization()
从上图可以看出，大约到20次的时候，惩罚效果就出来了。模型虽然很自信的能说出正确答案，但是给他一个小的惩罚，越自信，损失反而越大。
第一个例子
我们先从一个最简单的例子开始：从一个很小的词典中随机选取一些词做为输入，目标就是重新输出这些词。我们称为copy任务。
造数据
def data_gen(V, batch_size, nbatches):
"""
生成一组随机数据。（该方法仅用于Demo）
:param V: 词典的大小
:param batch_size
:param nbatches: 生成多少个batch
:return: yield一个Batch对象
"""
# 生成{nbatches}个batch
for i in range(nbatches):
# 生成一组输入数据
data = torch.randint(1, V, size=(batch_size, 10))
# 将每行的第一个词都改为1，即""
data[:, 0] = 1
# 该数据不需要梯度下降
src = data.requires_grad_(False).clone().detach()
tgt = data.requires_grad_(False).clone().detach()
# 返回一个Batch对象
yield Batch(src, tgt, 0)
损失计算
下面的类可不只是例子要用，最后的实战部分也要用：
class SimpleLossCompute:
"""
一个简单的损失计算和训练函数。
该类除了包含损失计算外，还包含模型generator部分的前向传递。
如果你对上面这句话不太理解，可参考这篇文章：
https://blog.csdn.net/zhaohongfei_358/article/details/125759911
请参考章节：Pytorch 实现梯度下降与参数更新
"""
def __init__(self, generator, criterion):
"""
generator: Generator类对象，用于根据Decoder的输出预测下一个token
criterion: LabelSmoothing类对象，用于对Label进行平滑和计算损失
"""
self.generator = generator
self.criterion = criterion
def __call__(self, x, y, norm):
"""
x: EncoderDecoder的输出，也就是Decoder的输出
y: batch.tgt_y，要被预测的所有token，例如src为` I love you `，
则`tgt_y`则为`我 爱 你 `
norm: batch.ntokens, tgt_y中的有效token数。用于对loss进行正则化。
"""
# 调用generator预测token。（EncoderDecoder的forward中并没有调用generator）
x = self.generator(x)
"""
这里首先使用KLDivLoss进行了损失计算，随后又除以batch.ntokens对loss进行正则化。
"""
sloss = (
self.criterion(
x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)
)
/ norm
)
return sloss.data * norm, sloss
这里详细说明一下上面的正则化，直接通过例子吧。
假设我们的第一个batch（假设batch_size=2）的tgt_y对应的句子为：
则我们的loss相当于对8个预测结果进行了计算，即我 爱 你 ，也就是 l o s s = ∑ i = 1 8 L ( t o k e n i ) loss = \sum_{i=1}^{8}L(token_i) loss=∑i=18L(tokeni)
假设我们的第二个batch的tgt_y对应的句子为：
则我们的loss岘港与对200个预测结果进行了计算，也就是 l o s s = ∑ i = 1 200 L ( t o k e n i ) loss = \sum_{i=1}^{200}L(token_i) loss=∑i=1200L(tokeni)
如果不做平均的话，显然第二个大，然后就会出现loss一会大一会小的情况，所以要求一下平均，也就是除以有效的token数。
使用贪心算法解码(Greedy Decoding)
简单起见，使用贪心算法解码(Greedy Decoding)进行翻译任务的预测。
在这里，所谓的贪心算法就是Transformer的正常推理过程：先求出encoder的输出memory，然后利用memory一个一个的求出token。
# 这个代码其实和最开始的inference_test()是一样的
def greedy_decode(model, src, src_mask, max_len, start_symbol):
"""
进行模型推理，推理出所有预测结果。
:param model: Transformer模型，即EncoderDecoder类对象
:param src: Encoder的输入inputs，Shape为(batch_size, 词数)
例如：[[1, 2, 3, 4, 5, 6, 7, 8, 0, 0]]
即一个句子，该句子有10个词，分别为1,2,...,0
:param src_mask: src的掩码，掩盖住非句子成分。
:param max_len: 一个句子的最大长度。
:param start_symbol: '' 对应的index，在本例中始终为0
:return: 预测结果，例如[[1, 2, 3, 4, 5, 6, 7, 8]]
"""
# 将src送入Transformer的Encoder，输出memory
memory = model.encode(src, src_mask)
# 初始化ys为[[0]]，用于保存预测结果，其中0表示''
ys = torch.zeros(1, 1).fill_(start_symbol).type_as(src.data)
# 循环调用decoder，一个个的进行预测。例如：假设我们要将“I love you”翻译成
# “我爱你”，则第一次的`ys`为()，然后输出为“I”。然后第二次`ys`为(, I)
# 输出为"love"，依次类推，直到decoder输出“”或达到句子最大长度。
for i in range(max_len - 1):
# 将encoder的输出memory和之前Decoder的所有输出作为参数，让Decoder来预测下一个token
out = model.decode(
# ys就是Decoder之前的所有输出
memory, src_mask, ys, subsequent_mask(ys.size(1)).type_as(src.data)
)
# 将Decoder的输出送给generator进行预测。这里只取最后一个词的输出进行预测。
# 因为你传的tgt的词数是变化的，第一次是()，第二次是(, I)
# 所以你的out的维度也是变化的，变化的就是(batch_size, 词数，词向量)中词数这个维度
# 前面的词向量送给generator预测的话，预测出来的也是前面的词，所以只能取最后一个。
prob = model.generator(out[:, -1])
# 取出数值最大的那个，它的index在词典中对应的词就是预测结果
_, next_word = torch.max(prob, dim=1)
# 取出预测结果
next_word = next_word.data[0]
# 将这一次的预测结果和之前的拼到一块，作为之后Decoder的输入
ys = torch.cat(
[ys, torch.zeros(1, 1).type_as(src.data).fill_(next_word)], dim=1
)
# 返回最终的预测结果。
return ys
接下来进行一个简单的模型训练，来实现copy任务。
# 训练一个简单的copy任务
def example_simple_model():
# 定义词典大小为11
V = 11
# 定义损失函数
criterion = LabelSmoothing(size=V, padding_idx=0, smoothing=0.0)
# 构建模型，src和tgt的词典大小都为2，Layer数量为2
model = make_model(V, V, N=2)
# 使用Adam优化器
optimizer = torch.optim.Adam(
model.parameters(), lr=0.5, betas=(0.9, 0.98), eps=1e-9
)
# 自定义Warmup学习率
lr_scheduler = LambdaLR(
optimizer=optimizer,
lr_lambda=lambda step: rate(
step, model_size=model.src_embed[0].d_model, factor=1.0, warmup=400
),
)
batch_size = 80
# 运行20个epoch
for epoch in range(20):
# 将模型调整为训练模式
model.train()
# 训练一个Batch
run_epoch(
# 生成20个batch对象
data_gen(V, batch_size, 20),
model,
SimpleLossCompute(model.generator, criterion),
optimizer,
lr_scheduler,
mode="train",
)
model.eval()
# 在一个epoch后，进行模型验证
run_epoch(
data_gen(V, batch_size, 5),
model,
SimpleLossCompute(model.generator, criterion),
# 不进行参数更新
DummyOptimizer(),
# 不调整学习率
DummyScheduler(),
mode="eval",
)[0] # run_epoch返回loss和train_state，这里取loss，所以是[0]
# 但是源码中并没有接收这个loss，所以这个[0]实际上没什么意义
# 将模型调整为测试模式，准备开始推理
model.eval()
# 定义一个src 0-9，看看model能不能重新输出0-9
src = torch.LongTensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
# 句子最大长度就是10,
max_len = src.shape[1]
# 不需要掩码，因为这10个都是有意义的数字
src_mask = torch.ones(1, 1, max_len)
# 使用greedy_decoder函数进行推理
print(greedy_decode(model, src, src_mask, max_len=max_len, start_symbol=0))
example_simple_model()
Epoch Step: 1 | Accumulation Step: 2 | Loss: 3.15 | Tokens / Sec: 822.1 | Learning Rate: 5.5e-06
Epoch Step: 1 | Accumulation Step: 2 | Loss: 2.05 | Tokens / Sec: 862.9 | Learning Rate: 6.1e-05
....略
Epoch Step: 1 | Accumulation Step: 2 | Loss: 0.12 | Tokens / Sec: 996.8 | Learning Rate: 1.0e-03
Epoch Step: 1 | Accumulation Step: 2 | Loss: 0.08 | Tokens / Sec: 999.6 | Learning Rate: 1.1e-03
tensor([[0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
如果运行了这个例子，你可能会有一点小疑问：①为什么Epoch Step始终为1；②为什么Accumulation Step始终为2
- 关于第一点：首先，Epoch Step指的是在本次epoch中执行到了第几个batch(从0开始)。其次因为在run_epoch中，会在第1,41,81,…个batch打印日志，而在本例中一共就20个batch，所以一个epoch只会打印一次日志。
- 第二点：Accumulation Step指的是在本次epoch中执行参数更新(
optimizier.step())的次数。在第一打印日志时，一共执行了2次参数更新，又因为一个epoch只会打印一次日志，所以该值始终为2。