base model第七弹：pytorch使用warm up、consine learning rate、label smooth、apex混合精度训练、梯度累加
文章目录
- warm up与consine learning rate
- label smooth
- apex混合精度训练
- 梯度累加
所有代码已上传到本人github repository：https://github.com/zgcr/pytorch-ImageNet-CIFAR-COCO-VOC-training
如果觉得有用，请点个star哟！
下列代码均在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%。在同样batchsize的情况下，使用apex会比不使用apex的情况下训练速度会稍稍变慢，但由于使用apex后，显存占用减少，因此可以继续扩大batchsize，这样相比原来的Batchsize训练速度还是要快很多。最后，在同样的batchsize情况下，使用apex训练出的模型性能和未使用apex时基本一致。
梯度累加
当我们的显卡资源有限时，训练时往往跑不起大的batchsize。梯度累加可以用较少的显卡模拟较大batchsize情况下的训练。具体来说，我们每次用小的batchsize计算出各个变量的梯度，然后把梯度累加起来，最后当小的batchsize累加和等于我们设定的大batchsize时再更新梯度。注意这个模拟较大Batchsize是近似的，因为bn层仍然按照小的batchsize来更新。
在train.py中，进行下列修改，就可以实现梯度累加：
inputs, labels = inputs.cuda(), labels.cuda()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss = loss / args.accumulation_steps
loss.backward()
if iter_index % args.accumulation_steps == 0:
optimizer.step()
optimizer.zero_grad()
其中，loss = loss / args.accumulation_steps根据累加次数把loss做除法，使得在不修改学习率的情况下最终累加的反传梯度数量级与原来一致（由于我们在累加的过程中考虑了多次小batchsize产生的loss，但又除以了累加次数，所以在loss.backward()多次累加后就近似相当于大batchsize计算出的loss产生的梯度）。最后，我们经过累加的梯度次数后才使用optimizer.step()更新一次网络参数，然后使用optimizer.zero_grad()将梯度清零。
需要注意的是，由于batchsize仍然按照真实的小batchsize更新，因此真实batchsize不宜设置的太小，否则误差会比较大。对于分类任务，一般真实batchsize不宜小于32。