base model第七弹:pytorch使用warm up、consine learning rate、label smooth、apex混合精度训练

下列代码均在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时基本一致。

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章