基于pytorch的自动混合精度(AMP)

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

1、什么是自动混合精度训练?

2、为什么需要自动混合精度?

3、如何在PyTorch中使用自动混合精度?

Pytorch 1.6版本以后,Pytorch将amp的功能吸收入官方库,位于torch.cuda.amp模块下。

【介绍】

torch.cuda.amp提供了对混合精度的支持。为实现自动混合精度训练,需要结合使用如下两个模块:

torch.cuda.amp.autocast:autocast主要用作上下文管理器或者装饰器,来确定使用混合精度的范围。

torch.cuda.amp.GradScalar:GradScalar主要用来完成梯度缩放。

【张量】

1、它们是专门用于特定类型矩阵运算的专用内核。


可以将两个FP16矩阵相乘并将其添加到FP16 / FP32矩阵中,从而得到FP16 / FP32矩阵。Tensor内核支持混合精度数学,即输入为半精度(FP16),输出为全精度(FP32)。上面的操作对于许多深度学习任务具有内在的价值,并且Tensor内核为该操作提供了专用的硬件。

现在,使用FP16和FP32主要有两个好处

FP16需要较少的内存,因此更易于训练和部署大型神经网络。它还减少了数据移动。

使用Tensor Core,数学运算的运行速度大大降低了精度。NVIDIA提供的Volta GPU的确切数量是:FP16中为125 TFlops,而FP32中为15.7 TFlops(加速8倍)

但是也有缺点。从FP32转到FP16时,必然会降低精度。


FP32与FP16:FP32具有八个指数位和23个小数位,而FP16具有五个指数位和十个小数位。

但是需要FP32吗?

FP16实际上可以很好地表示大多数权重和渐变。因此,拥有存储和使用FP32所需的所有这些额外位只是浪费。

如何使用Tensor Core】

NVIDIA可以轻松地将Tensor内核与自动混合精度一起使用,并提供了几行代码。需要在代码中做两件事:

1.FP32所需的操作(如Softmax)被分配给FP32,而FP16可以完成的操作(如Conv)被自动分配给FP16。

2.使用损耗定标保留较小的梯度值。梯度值可能超出FP16的范围。在这种情况下,将对梯度值进行缩放,使其落在FP16范围内。


pytorch从1.6版本开始,已经内置了torch.cuda.amp,采用自动混合精度训练就不需要加载第三方NVIDIA的apex库了。

【优点】

   1.减少显存占用;

 2.加快训练和推断的计算,能带来多一倍速的体验;

 3.张量核心的普及(NVIDIA Tensor Core),低精度计算是未来深度学习的一个重要趋势。

  但凡事都有两面性,FP16也带来了些问题:1.溢出错误;2.舍入误差;

1.溢出错误:由于FP16的动态范围比FP32位的狭窄很多,因此,在计算过程中很容易出现上溢出和下溢出,溢出之后就会出现"NaN"的问题。在深度学习中,由于激活函数的梯度往往要比权重梯度小,更易出现下溢出的情况。


2.舍入误差

 舍入误差指的是当梯度过小时,小于当前区间内的最小间隔时,该次梯度更新可能会失败:

为了消除torch.HalfTensor也就是FP16的问题,需要使用以下两种方法:

1)混合精度训练

 在内存中用FP16做储存和乘法从而加速计算,而用FP32做累加避免舍入误差。混合精度训练的策略有效地缓解了舍入误差的问题。

什么时候用torch.FloatTensor,什么时候用torch.HalfTensor呢?这是由pytorch框架决定的,在pytorch1.6的AMP上下文中,以下操作中Tensor会被自动转化为半精度浮点型torch.HalfTensor:

__matmul__

addbmm

addmm

addmv

addr

baddbmm

bmm

chain_matmul

conv1d

conv2d

conv3d

conv_transpose1d

conv_transpose2d

conv_transpose3d

linear

matmul

mm

mv

prelu

2)损失放大(Loss scaling)

即使了混合精度训练,还是存在无法收敛的情况,原因是激活梯度的值太小,造成了溢出。可以通过使用torch.cuda.amp.GradScaler,通过放大loss的值来防止梯度的underflow(只在BP时传递梯度信息使用,真正更新权重时还是要把放大的梯度再unscale回去);

反向传播前,将损失变化手动增大2^k倍,因此反向传播时得到的中间变量(激活函数梯度)则不会溢出;

反向传播后,将权重梯度缩小2^k倍,恢复正常值。

【使用】

pytorch1.6及以上版本

有两个接口:autocast和Gradscaler

1) autocast

导入pytorch中模块torch.cuda.amp的类autocast

from torch.cuda.amp import autocast as autocast

model=Net().cuda()

optimizer=optim.SGD(model.parameters(),...)

for input,target in data:

optimizer.zero_grad()

with autocast():

output=model(input)

loss = loss_fn(output,target)

loss.backward()

optimizer.step()

可以使用autocast的context managers语义(如上),也可以使用decorators语义。当进入autocast上下文后,在这之后的cuda ops会把tensor的数据类型转换为半精度浮点型,从而在不损失训练精度的情况下加快运算。而不需要手动调用.half(),框架会自动完成转换。

不过,autocast上下文只能包含网络的前向过程(包括loss的计算),不能包含反向传播,因为BP的op会使用和前向op相同的类型。

2)GradScaler

  使用前,需要在训练最开始前实例化一个GradScaler对象,例程如下:

from torch.cuda.amp import autocast as autocast

model=Net().cuda()

optimizer=optim.SGD(model.parameters(),...)

scaler = GradScaler() #训练前实例化一个GradScaler对象

for epoch in epochs:

for input,target in data:

optimizer.zero_grad()

with autocast(): #前后开启autocast

output=model(input)

loss = loss_fn(output,targt)

scaler.scale(loss).backward() #为了梯度放大

#scaler.step() 首先把梯度值unscale回来,如果梯度值不是inf或NaN,则调用optimizer.step()来更新权重,否则,忽略step调用,从而保证权重不更新。

scaler.step(optimizer)

scaler.update() #准备着,看是否要增大scaler。

scaler的大小在每次迭代中动态估计,为了尽可能减少梯度underflow,scaler应该更大;但太大,半精度浮点型又容易overflow(变成inf或NaN).所以,动态估计原理就是在不出现if或NaN梯度的情况下,尽可能的增大scaler值。在每次scaler.step(optimizer)中,都会检查是否有inf或NaN的梯度出现:

       1.如果出现inf或NaN,scaler.step(optimizer)会忽略此次权重更新(optimizer.step()),并将scaler的大小缩小(乘上backoff_factor);

  2.如果没有出现inf或NaN,那么权重正常更新,并且当连续多次(growth_interval指定)没有出现inf或NaN,则scaler.update()会将scaler的大小增加(乘上growth_factor)。

对于分布式训练,由于autocast是thread local的,要注意以下情形:

1)torch.nn.DataParallel:

以下代码分布式是不生效的

model = MyModel()

dp_model = nn.DataParallel(model)

with autocast():

output=dp_model(input)

loss=loss_fn(output)

需使用autocast装饰model的forward函数


2)torch.nn.DistributedDataParallel:

同样,对于多GPU,也需要autocast装饰model的forward方法,保证autocast在进程内部生效。

【注意】

1.判断GPU是否支持FP16,支持Tensor core的GPU(2080Ti,Titan,Tesla等),不支持的(Pascal系列)不建议;

2.常数范围:为了保证计算不溢出,首先保证人工设定的常数不溢出。如epsilon,INF等;

3.Dimension最好是8的倍数:维度是8的倍数,性能最好;

4.涉及sum的操作要小心,容易溢出,softmax操作,建议用官方API,并定义成layer写在模型初始化里;

5.模型书写要规范:自定义的Layer写在模型初始化函数里,graph计算写在forward里;

6.一些不常用的函数,使用前要注册:amp.register_float_function(torch,'sogmoid')

7.某些函数不支持FP16加速,建议不要用;

8.需要操作梯度的模块必须在optimizer的step里,不然AMP不能判断grad是否为NaN。

torch.HalfTensor的优势就是存储小、计算快、更好的利用CUDA设备的Tensor Core。因此训练的时候可以减少显存的占用(可以增加batchsize了),同时训练速度更快;

torch.HalfTensor的劣势就是:数值范围小(更容易Overflow / Underflow)、舍入误差(Rounding Error,导致一些微小的梯度信息达不到16bit精度的最低分辨率,从而丢失)。

可见,当有优势的时候就用torch.HalfTensor,而为了消除torch.HalfTensor的劣势,我们带来了两种解决方案:

1,梯度scale,这正是上一小节中提到的torch.cuda.amp.GradScaler,通过放大loss的值来防止梯度的underflow(这只是BP的时候传递梯度信息使用,真正更新权重的时候还是要把放大的梯度再unscale回去);

2,回落到torch.FloatTensor,这就是混合一词的由来。那怎么知道什么时候用torch.FloatTensor,什么时候用半精度浮点型呢?这是PyTorch框架决定的。

Working with Multiple GPUs】

针对多卡训练的情况,只影响autocast的使用方法,GradScaler的用法与之前一致。

.1 DataParallel in a single process

有效的调用方式如下所示:


【 DistributedDataParallel, multiple GPUs per process】

与DataParallel的使用相同,在模型构建时,对forward函数的定义方式进行修改,保证autocast在进程内部生效。

【参考文章:https://zhuanlan.zhihu.com/p/408610877】

©著作权归作者所有,转载或内容合作请联系作者
  • 序言:七十年代末,一起剥皮案震惊了整个滨河市,随后出现的几起案子,更是在滨河造成了极大的恐慌,老刑警刘岩,带你破解...
    沈念sama阅读 212,332评论 6 493
  • 序言:滨河连续发生了三起死亡事件,死亡现场离奇诡异,居然都是意外死亡,警方通过查阅死者的电脑和手机,发现死者居然都...
    沈念sama阅读 90,508评论 3 385
  • 文/潘晓璐 我一进店门,熙熙楼的掌柜王于贵愁眉苦脸地迎上来,“玉大人,你说我怎么就摊上这事。” “怎么了?”我有些...
    开封第一讲书人阅读 157,812评论 0 348
  • 文/不坏的土叔 我叫张陵,是天一观的道长。 经常有香客问我,道长,这世上最难降的妖魔是什么? 我笑而不...
    开封第一讲书人阅读 56,607评论 1 284
  • 正文 为了忘掉前任,我火速办了婚礼,结果婚礼上,老公的妹妹穿的比我还像新娘。我一直安慰自己,他们只是感情好,可当我...
    茶点故事阅读 65,728评论 6 386
  • 文/花漫 我一把揭开白布。 她就那样静静地躺着,像睡着了一般。 火红的嫁衣衬着肌肤如雪。 梳的纹丝不乱的头发上,一...
    开封第一讲书人阅读 49,919评论 1 290
  • 那天,我揣着相机与录音,去河边找鬼。 笑死,一个胖子当着我的面吹牛,可吹牛的内容都是我干的。 我是一名探鬼主播,决...
    沈念sama阅读 39,071评论 3 410
  • 文/苍兰香墨 我猛地睁开眼,长吁一口气:“原来是场噩梦啊……” “哼!你这毒妇竟也来了?” 一声冷哼从身侧响起,我...
    开封第一讲书人阅读 37,802评论 0 268
  • 序言:老挝万荣一对情侣失踪,失踪者是张志新(化名)和其女友刘颖,没想到半个月后,有当地人在树林里发现了一具尸体,经...
    沈念sama阅读 44,256评论 1 303
  • 正文 独居荒郊野岭守林人离奇死亡,尸身上长有42处带血的脓包…… 初始之章·张勋 以下内容为张勋视角 年9月15日...
    茶点故事阅读 36,576评论 2 327
  • 正文 我和宋清朗相恋三年,在试婚纱的时候发现自己被绿了。 大学时的朋友给我发了我未婚夫和他白月光在一起吃饭的照片。...
    茶点故事阅读 38,712评论 1 341
  • 序言:一个原本活蹦乱跳的男人离奇死亡,死状恐怖,灵堂内的尸体忽然破棺而出,到底是诈尸还是另有隐情,我是刑警宁泽,带...
    沈念sama阅读 34,389评论 4 332
  • 正文 年R本政府宣布,位于F岛的核电站,受9级特大地震影响,放射性物质发生泄漏。R本人自食恶果不足惜,却给世界环境...
    茶点故事阅读 40,032评论 3 316
  • 文/蒙蒙 一、第九天 我趴在偏房一处隐蔽的房顶上张望。 院中可真热闹,春花似锦、人声如沸。这庄子的主人今日做“春日...
    开封第一讲书人阅读 30,798评论 0 21
  • 文/苍兰香墨 我抬头看了看天上的太阳。三九已至,却和暖如春,着一层夹袄步出监牢的瞬间,已是汗流浃背。 一阵脚步声响...
    开封第一讲书人阅读 32,026评论 1 266
  • 我被黑心中介骗来泰国打工, 没想到刚下飞机就差点儿被人妖公主榨干…… 1. 我叫王不留,地道东北人。 一个月前我还...
    沈念sama阅读 46,473评论 2 360
  • 正文 我出身青楼,却偏偏与公主长得像,于是被迫代替她去往敌国和亲。 传闻我的和亲对象是个残疾皇子,可洞房花烛夜当晚...
    茶点故事阅读 43,606评论 2 350

推荐阅读更多精彩内容