基于 PyTorch 的混合精度训练加速

姓名:唐易平

学号:19121110448

转载自:【PyTorch】唯快不破:基于Apex的混合精度加速

【嵌牛导读】:介绍一下混合精度计算(Mixed Precision),并分享一款Nvidia开发的基于PyTorch的混合精度训练加速神器--Apex。

【嵌牛鼻子】:PyTorch,混合精度训练

【嵌牛提问】:Apex加速原理?如何使用?

【嵌牛正文】:

理论部分:

为了充分理解混合精度的原理,以及API的使用,先补充一点基础的理论知识。

1. 什么是FP16?

半精度浮点数是一种计算机使用的二进制浮点数数据类型,使用2字节(16位)存储。

FP16和FP32表示的范围和精度对比  

其中,sign位表示正负,exponent位表示指数(2^{n-15+1(n=0)}),fraction表示的是分数(\frac{m}{1024} ),其中当指数为零的时候,下图加号左边为0,其他情况为1。

FP16的表示范例

2. 为什么需要FP16?

1). 减少显存占用

现在模型越来越大,当你使用Bert这一类的预训练模型时,往往显存就被模型及模型计算占去大半,当想要使用更大的Batch Size的时候会显得捉襟见肘。由于FP16的内存占用只有FP32的一半,自然地就可以帮助训练过程节省一半的显存空间。

2). 加快训练和推断的计算

与普通的空间时间Trade-off的加速方法不同,FP16除了能节约内存,还能同时节省模型的训练时间。在大部分的测试中,基于FP16的加速方法能够给模型训练带来多一倍的加速体验。

3). 张量核心的普及

硬件的发展同样也推动着模型计算的加速,随着Nvidia张量核心(Tensor Core)的普及,16bit计算也一步步走向成熟,低精度计算也是未来深度学习的一个重要趋势,再不学习就out啦。

3. FP16带来的问题:量化误差

1). 溢出错误(Grad Overflow / Underflow)

由于FP16的动态范围(6×10^{-8}\sim65504 )比FP32的动态范围(1.4\times10^{-45}\sim1.7\times10^{38})要狭窄很多,因此在计算过程中很容易出现上溢出(Overflow,g>65504 )和下溢出(Underflow,g<6\times10^{-8})的错误,溢出之后就会出现“Nan”的问题。

在深度学习中,由于激活函数的的梯度往往要比权重梯度小,更易出现下溢出的情况。

下溢出问题

2). 舍入误差(Rounding Error)

舍入误差指的是当梯度过小,小于当前区间内的最小间隔时,该次梯度更新可能会失败,用一张图清晰地表示:

舍入误差

4. 解决问题的办法:混合精度训练+动态损失放大

1). 混合精度训练(Mixed Precision)

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

2). 损失放大(Loss Scaling)

即使用了混合精度训练,还是会存在无法收敛的情况,原因是激活梯度的值太小,造成了下溢出(Underflow)。损失放大的思路是:

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

\cdot 反向传播后,将权重梯度缩2k2k倍,恢复正常值。


Apex的新API:Automatic Mixed Precision (AMP)

曾经的Apex混合精度训练的api仍然需要手动half模型已经输入的数据,比较麻烦,现在新的api只需要三行代码即可无痛使用:

from apex import amp

model, optimizer = amp.initialize(model, optimizer, opt_level="O1")# 这里是“欧一”,不是“零一”

with amp.scale_loss(loss, optimizer) as scaled_loss:

         scaled_loss.backward()

1.opt_level

其中只有一个opt_level需要用户自行配置:

\cdot  O0:纯FP32训练,可以作为accuracy的baseline;

\cdot  O1:混合精度训练(推荐使用),根据黑白名单自动决定使用FP16(GEMM, 卷积)还是FP32(Softmax)进行计算。

\cdot  O2:“几乎FP16”混合精度训练,不存在黑白名单,除了Batch norm,几乎都是用FP16计算。

\cdot  O3:纯FP16训练,很不稳定,但是可以作为speed的baseline;

2.动态损失放大(Dynamic Loss Scaling)

AMP默认使用动态损失放大,为了充分利用FP16的范围,缓解舍入误差,尽量使用最高的放大倍数(2^{24}),如果产生了上溢出(Overflow),则跳过参数更新,缩小放大倍数使其不溢出,在一定步数后(比如2000步)会再尝试使用大的scale来充分利用FP16的范围:

AMP中动态损失放大的策略

干货:踩过的那些坑

\cdot  判断你的GPU是否支持FP16:构拥有Tensor Core的GPU(2080Ti、Titan、Tesla等),不支持的(Pascal系列)就不建议折腾了。

\cdot  常数的范围:为了保证计算不溢出,首先要保证人为设定的常数(包括调用的源码中的)不溢出,如各种epsilon,INF等。

\cdot  Dimension最好是8的倍数:Nvidia官方文档表示,维度都是8的倍数的时候,性能最好。

\cdot  涉及到sum的操作要小心,很容易溢出,类似Softmax的操作建议用官方API,并定义成layer写在模型初始化里。

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

\cdot  某些不常用的函数,在使用前需要注册:amp.register_float_function(torch, 'sigmoid')

\cdot  某些函数(如einsum)暂不支持FP16加速,建议不要用的太heavy。

\cdot  需要操作模型参数的模块(类似EMA),要使用AMP封装后的model。

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

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

推荐阅读更多精彩内容