pytorch 如何使用amp进行混合精度训练


Posted in Python onMay 24, 2021

简介

AMP:Automatic mixed precision,自动混合精度,可以在神经网络推理过程中,针对不同的层,采用不同的数据精度进行计算,从而实现节省显存和加快速度的目的。

在Pytorch 1.5版本及以前,通过NVIDIA提供的apex库可以实现amp功能。但是在使用过程中会伴随着一些版本兼容和奇怪的报错问题。

从1.6版本开始,Pytorch原生支持自动混合精度训练,并已进入稳定阶段,AMP 训练能在 Tensor Core GPU 上实现更高的性能并节省多达 50% 的内存。

环境

Python 3.8

Pytorch 1.7.1

CUDA 11 + cudnn 8

NVIDIA GeFore RTX 3070

ps:后续使用移动端的3070,或者3080结合我目前训练的分类网络来测试实际效果

原理

关于低精度计算

当前的深度学习框架大都采用的都是FP32来进行权重参数的存储,比如Python float的类型为双精度浮点数 FP64,PyTorch Tensor的默认类型为单精度浮点数FP32。

随着模型越来越大,加速训练模型的需求就产生了。在深度学习模型中使用FP32主要存在几个问题,第一模型尺寸大,训练的时候对显卡的显存要求高;第二模型训练速度慢;第三模型推理速度慢。

其解决方案就是使用低精度计算对模型进行优化。

推理过程中的模型优化目前比较成熟的方案就是FP16量化和INT8量化,NVIDIA TensorRT等框架就可以支持,这里不再赘述。训练方面的方案就是混合精度训练,它的基本思想很简单: 精度减半(FP32→ FP16) ,训练时间减半。

与单精度浮点数float32(32bit,4个字节)相比,半精度浮点数float16仅有16bit,2个字节组成。

可以很明显的看到,使用FP16可以解决或者缓解上面FP32的两个问题:显存占用更少:通用的模型FP16占用的内存只需原来的一半,训练的时候可以使用更大的batchsize。

计算速度更快:有论文指出半精度的计算吞吐量可以是单精度的 2-8 倍。

pytorch 如何使用amp进行混合精度训练

从上到下依次为 fp16、fp32 、fp64

当前很多NVIDIA GPU搭载了专门为快速FP16矩阵运算设计的特殊用途Tensor Core,比如Tesla P100,Tesla V100、Tesla A100、GTX 20XX 和RTX 30XX等。

Tensor Core是一种矩阵乘累加的计算单元,每个Tensor Core每个时钟执行64个浮点混合精度操作(FP16矩阵相乘和FP32累加),英伟达宣称使用Tensor Core进行矩阵运算可以轻易的提速,同时降低一半的显存访问和存储。

随着Tensor Core的普及FP16计算也一步步走向成熟,低精度计算也是未来深度学习的一个重要趋势。

pytorch 如何使用amp进行混合精度训练

Tensor Core 的 4x4x4 矩阵乘法与累加

pytorch 如何使用amp进行混合精度训练

Volta GV100 Tensor Core 流程图

自动混合精度训练

不同于在推理过程中直接削减权重精度,在模型训练的过程中,直接使用半精度进行计算会导致的两个问题的处理:舍入误差(Rounding Error)和溢出错误(Grad Overflow / Underflow)。

舍入误差: float16的最大舍入误差约为 (~2 ^-10 ),比float32的最大舍入误差(~2 ^-23) 要大不少。 对足够小的浮点数执行的任何操作都会将该值四舍五入到零,在反向传播中很多甚至大多数梯度更新值都非常小,但不为零。 在反向传播中舍入误差累积可以把这些数字变成0或者 nan, 这会导致不准确的梯度更新,影响网络的收敛。

溢出错误: 由于float16的有效的动态范围约为 ( 5.96×10^-8 ~ 6.55×10^4),比单精度的float32(1.4x10^-45 ~ 1.7x10^38)要狭窄很多,精度下降(小数点后16相比较小数点后8位要精确的多)会导致得到的值大于或者小于fp16的有效动态范围,也就是上溢出或者下溢出。

在深度学习中,由于激活函数的的梯度往往要比权重梯度小,更易出现下溢出的情况。2018年ICLR论文 Mixed Precision Training 中提到,简单的在每个地方使用FP16会损失掉梯度更新小于2^-24的值——大约占他们的示例网络所有梯度更新的5%。

解决方案就是使用混合精度训练(Mixed Precision)和损失缩放(Loss Scaling):

1、混合精度训练:

混合精度训练是一种通过在FP16上执行尽可能多的操作来大幅度减少神经网络训练时间的技术,在像线性层或是卷积操作上,FP16运算较快,但像Reduction运算又需要 FP32的动态范围。通过混合精度训练的方式,便可以在部分运算操作使用FP16,另一部分则使用 FP32,混合精度功能会尝试为每个运算使用相匹配的数据类型,在内存中用FP16做储存和乘法从而加速计算,用FP32做累加避免舍入误差。这样在权重更新的时候就不会出现舍入误差导致更新失败,混合精度训练的策略有效地缓解了舍入误差的问题。

2、损失缩放:

即使用了混合精度训练,还是会存在无法收敛的情况,原因是激活梯度的值太小,造成了下溢出。损失缩放是指在执行反向传播之前,将损失函数的输出乘以某个标量数(论文建议从8开始)。 乘性增加的损失值产生乘性增加的梯度更新值,提升许多梯度更新值到超过FP16的安全阈值2^-24。 只要确保在应用梯度更新之前撤消缩放,并且不要选择一个太大的缩放以至于产生inf权重更新(上溢出) ,从而导致网络向相反的方向发散。

使用Pytorch AMP

Pytorch原生的amp模式使用起来相当简单,只需要从torch.cuda.amp导入GradScaler和 autocast这两个函数即可。torch.cuda.amp的名字意味着这个功能只能在cuda上使用,事实上,这个功能正是NVIDIA的开发人员贡献到PyTorch项目中的。

Pytorch在amp模式下维护两个权重矩阵的副本,一个主副本用 FP32,一个半精度副本用 FP16。 梯度更新使用FP16矩阵计算,但更新于 FP32矩阵。 这使得应用梯度更新更加安全。

autocast上下文管理器实现了 FP32到FP16的转换,它会自动判别哪些层可以进行FP16哪些层不可以。 GradScaler对梯度更新计算(检查是否溢出)和优化器(将丢弃的batches转换为 no-op)进行控制,通过放大loss的值来防止梯度的溢出。

在训练中的具体使用方法如下所示:

def train():
    batch_size = 8
    epochs = 10
    lr = 1e-3
    size = 256
    num_class = 35
    use_amp = True
 
    device = 'cuda' if torch.cuda.is_available() else 'cpu'
 
    print('torch version: {}'.format(torch.__version__))
    print('amp:           {}'.format(use_amp))
    print('device:        {}'.format(device))
    print('epochs:        {}'.format(epochs))
    print('learn rate:    {}'.format(lr))
    print('batch size:    {}'.format(batch_size))
 
    net = ERFNet(num_classes=num_class).to(device)
 
    train_data = CityScapesDataset('D:\\dataset\\cityscapes',
                                   'D:\\dataset\\cityscapes\\trainImages.txt',
                                   'D:\\dataset\\cityscapes\\trainLabels.txt',
                                   size, num_class)
    val_data = CityScapesDataset('D:\\dataset\\cityscapes',
                                 'D:\\dataset\\cityscapes\\valImages.txt',
                                 'D:\\dataset\\cityscapes\\valLabels.txt',
                                 size, num_class)
 
    train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=False, num_workers=8)
    val_dataloader = DataLoader(val_data, batch_size=batch_size, shuffle=False, num_workers=4)
 
    opt = torch.optim.Adam(net.parameters(), lr=lr)
    criterion = torch.nn.CrossEntropyLoss(ignore_index=255)
 
    if use_amp:
        scaler = torch.cuda.amp.GradScaler()
 
    writer = SummaryWriter("summary")
 
    train_loss = AverageMeter()
    val_acc = AverageMeter()
    val_miou = AverageMeter()
 
    for epoch in range(0, epochs):
        train_loss.reset()
        val_acc.reset()
        val_miou.reset()
 
        with tqdm(total=train_data.__len__(), unit='img', desc="Epoch {}/{}".format(epoch + 1, epochs)) as pbar:
            # train
            net.train()
            for img, mask in train_dataloader:
                img = img.to(device)
                mask = mask.to(device)
                n = img.size()[0]
 
                opt.zero_grad()
 
                if use_amp:
                    with torch.cuda.amp.autocast():
                        output = net(img)
                        loss = criterion(output, mask)
 
                    scaler.scale(loss).backward()
                    scaler.step(opt)
                    scaler.update()
                else:
                    output = net(img)
                    loss = criterion(output, mask)
                    loss.backward()
                    opt.step()
 
                train_loss.update(loss.item(), n)
 
                pbar.set_postfix(**{"loss": train_loss.avg})
                pbar.update(img.size()[0])
 
            writer.add_scalar('train_loss', train_loss.avg, epoch)
            # eval
            net.eval()
            for img, mask in val_dataloader:
                img = img.to(device)
                mask = mask
                n = img.size()[0]
 
                output = net(img)
 
                pred_mask = torch.softmax(output, dim=1)
                pred_mask = pred_mask.detach().cpu().numpy()
                pred_mask = np.argmax(pred_mask, axis=1)
                true_mask = mask.numpy()
                acc, acc_cls, mean_iu, fwavacc = evaluate(pred_mask, true_mask, num_class)
 
                val_acc.update(acc)
                val_miou.update(mean_iu)
 
            writer.add_scalar('val_acc', val_acc.avg, epoch)
            writer.add_scalar('val_miou', val_miou.avg, epoch)
 
            pbar.set_postfix(**{"loss": train_loss.avg, "val_acc": val_acc.avg, "val_miou": val_miou.avg})

实验

硬件使用NVIDIA Geforce RTX 3070作为测试卡,这块卡有184个Tensor Core,能比较好的支持amp模式。

模型使用ERFNet分割模型作为基准,cityscapes作为测试数据,10个epoch下的测试效果如下所示:

在模型的训练性能方面,amp模式下的平均训练时间并没有明显节省,甚至还略低于正常模式。

显存的占用大约节省了25%,对于需要大量显存的模型来说这个提升还是相当可观的。

理论上训练速度应该也是有提升的,到Pytorch的GitHub issue里翻了一下,好像30系显卡会存在速度提不上来的问题,不太清楚是驱动支持不到位还是软件适配不到位。

Metrics time memory
AMP 66.72s 2.5G
NO_AMP 65.64s 3.3G

pytorch 如何使用amp进行混合精度训练

amp

pytorch 如何使用amp进行混合精度训练

no_amp

在模型的精度方面,在不进行数据shuffle的情况下统计了10个epoch下两种模式的train_loss和val_acc,可以看出不管是训练还是推理,amp模式并没有带来明显的精度损失。

pytorch 如何使用amp进行混合精度训练

cmp

以上为个人经验,希望能给大家一个参考,也希望大家多多支持三水点靠木。

Python 相关文章推荐
新手该如何学python怎么学好python?
Oct 07 Python
Python新手在作用域方面经常容易碰到的问题
Apr 03 Python
python修改操作系统时间的方法
May 18 Python
Django框架使用mysql视图操作示例
May 15 Python
对python中的os.getpid()和os.fork()函数详解
Aug 08 Python
python集成开发环境配置(pycharm)
Feb 14 Python
Python稀疏矩阵及参数保存代码实现
Apr 18 Python
TensorFLow 数学运算的示例代码
Apr 21 Python
.img/.hdr格式转.nii格式的操作
Jul 01 Python
安装pyecharts1.8.0版本后导入pyecharts模块绘图时报错: “所有图表类型将在 v1.9.0 版本开始强制使用 ChartItem 进行数据项配置 ”的解决方法
Aug 18 Python
python利用google翻译方法实例(翻译字幕文件)
Sep 21 Python
python爬取企查查企业信息之selenium自动模拟登录企查查
Apr 08 Python
只需要这一行代码就能让python计算速度提高十倍
pytorch 如何使用float64训练
pytorch 使用半精度模型部署的操作
May 24 #Python
解决Pytorch半精度浮点型网络训练的问题
May 24 #Python
Python办公自动化之Excel(中)
May 24 #Python
PyTorch梯度裁剪避免训练loss nan的操作
May 24 #Python
python3读取文件指定行的三种方法
May 24 #Python
You might like
介绍几个array库的新函数 php
2006/12/29 PHP
PHP Memcached应用实现代码
2010/02/08 PHP
PHPWind与Discuz截取字符函数substrs与cutstr性能比较
2011/12/05 PHP
WordPress的主题编写中获取头部模板和底部模板
2015/12/28 PHP
详谈配置phpstorm完美支持Codeigniter(CI)代码自动完成(代码提示)
2017/04/07 PHP
PHP中的访问修饰符简单比较
2019/02/02 PHP
javascript JSON操作入门实例
2010/04/16 Javascript
有关JavaScript的10个怪癖和秘密分享
2011/08/28 Javascript
js获得地址栏?问号后参数的方法
2013/08/08 Javascript
jQuery aminate方法定位到页面具体位置
2013/12/26 Javascript
详解JavaScript数组的操作大全
2015/10/19 Javascript
利用Angularjs和Bootstrap前端开发案例实战
2016/08/27 Javascript
JS实现简单易用的手机端浮动窗口显示效果
2016/09/07 Javascript
jQuery实现磁力图片跟随效果完整示例
2016/09/16 Javascript
jquery.multiselect多选下拉框实现代码
2016/11/11 Javascript
js下拉菜单生成器dropMenu使用方法详解
2017/08/01 Javascript
vuejs使用$emit和$on进行组件之间的传值的示例
2017/10/04 Javascript
Node.js API详解之 module模块用法实例分析
2020/05/13 Javascript
JavaScript this关键字指向常用情况解析
2020/09/02 Javascript
Python中多线程及程序锁浅析
2015/01/21 Python
在Django中进行用户注册和邮箱验证的方法
2016/05/09 Python
Python编程之序列操作实例详解
2017/07/22 Python
基于Django用户认证系统详解
2018/02/21 Python
对Python中class和instance以及self的用法详解
2019/06/26 Python
python多进程下实现日志记录按时间分割
2019/07/22 Python
jenkins配置python脚本定时任务过程图解
2019/10/29 Python
Python urllib库如何添加headers过程解析
2020/10/05 Python
HTML5 声明兼容IE的写法
2011/05/16 HTML / CSS
英国汽车和货车租赁网站:Hertz英国
2016/09/02 全球购物
Evisu官方网站:日本牛仔品牌,时尚街头设计风格
2016/12/30 全球购物
Happy Plugs官网:瑞典无线耳机品牌
2020/07/16 全球购物
零件设计自荐信范文
2013/11/27 职场文书
教师职称自我鉴定
2014/02/12 职场文书
2014年预备党员学习新党章思想汇报
2014/09/15 职场文书
试用期转正后的自我评价
2014/09/21 职场文书
2014小学生国庆65周年演讲稿
2014/09/21 职场文书