下列代码均在pytorch1.4版本中测试过,确认正确无误。
warm up与consine learning rate
warm up最早来自于这篇文章:https://arxiv.org/pdf/1706.02677.pdf 。根据这篇文章,我们一般只在前5个epoch使用warm up。consine learning rate来自于这篇文章:https://arxiv.org/pdf/1812.01187.pdf 。通常情况下,把warm up和consine learning rate一起使用会达到更好的效果。
代码实现:
# MultiStepLR without warm up
scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, milestones=args.milestones, gamma=0.1)
# warm_up_with_multistep_lr
warm_up_with_multistep_lr = lambda epoch: epoch / args.warm_up_epochs if epoch <= args.warm_up_epochs else 0.1**len([m for m in args.milestones if m <= epoch])
scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=warm_up_with_multistep_lr)
# warm_up_with_cosine_lr
warm_up_with_cosine_lr = lambda epoch: epoch / args.warm_up_epochs if epoch <= args.warm_up_epochs else 0.5 * ( math.cos((epoch - args.warm_up_epochs) /(args.epochs - args.warm_up_epochs) * math.pi) + 1)
scheduler = torch.optim.lr_scheduler.LambdaLR( optimizer, lr_lambda=warm_up_with_cosine_lr)
上面的三段代码分别是不使用warm up+multistep learning rate 衰减、使用warm up+multistep learning rate 衰减、使用warm up+consine learning rate衰减。代码均使用pytorch中的lr_scheduler。LambdaLR自定义学习率衰减器。
label smooth
对于分类任务,通常我们使用cross entropy损失函数进行训练,即:
criterion = nn.CrossEntropyLoss().cuda()
在这篇文章:https://arxiv.org/pdf/1812.01187.pdf 中给出了带label smooth的cross entropy函数计算公式。我们在pytorch中实现了它,代码如下:
class LabelSmoothCELoss(nn.Module):
def __init__(self):
super().__init__()
def forward(self, pred, label, smoothing=0.1):
pred = F.softmax(pred, dim=1)
one_hot_label = F.one_hot(label, pred.size(1)).float()
smoothed_one_hot_label = (
1.0 - smoothing) * one_hot_label + smoothing / pred.size(1)
loss = (-torch.log(pred)) * smoothed_one_hot_label
loss = loss.sum(axis=1, keepdim=False)
loss = loss.mean()
return loss
使用时,只需要:
criterion = LabelSmoothCELoss().cuda()
替换掉原来的cross entropy损失函数即可。
apex混合精度训练
apex是NVIDIA发布的开源混合精度训练工具库。该工具库提供了AMP(自动混合精度)和FP16_Optimizer两种不同的库。AMP提供自动混合精度训练支持,该库会自动检查执行的操作,对于FP16安全的操作在训练中Cast到FP16精度,反之则选择FP32精度。FP16_Optimizer提供的是高阶版混合精度训练,其提供了更多详细的实现细节,对于整个网络完全采用FP16精度训练。
这里只介绍如何用apex实现最简单的自动混合精度训练支持。
首先安装apex工具库。如果你是Python3.7环境,使用下列命令安装:
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
如果你是Python3.6环境,则使用下列命令安装:
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir ./
安装完成后,在train.py中进行下列修改:
from apex import amp
......
......
# 定义好Model和optimizer后,增加下面这行代码:
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
# 反传梯度原本的代码是loss.backward(),改为下面两行代码:
with amp.scale_loss(loss, optimizer) as scaled_loss:
scaled_loss.backward()
这样在训练时就会使用apex自动混合精度训练了。使用apex后,对于同一个训练任务,可有效降低显存占用25-30%,但训练速度会稍稍变慢。最后训练出的模型性能和未使用apex时基本一致。