[[313508]]

百度和Nvidia研究院結(jié)合N卡底層計算優(yōu)化，提出了一種有效的神經(jīng)網(wǎng)絡訓練加速方法，不僅是預訓練，在全民finetune BERT的今天變得異常有用。

一切還要從2018年ICLR的一篇論文說起。

《MIXED PRECISION TRAINING》是百度&Nvidia研究院一起發(fā)表的，結(jié)合N卡底層計算優(yōu)化，提出了一種灰常有效的神經(jīng)網(wǎng)絡訓練加速方法，不僅是預訓練，在全民finetune BERT的今天變得異常有用哇。

而且調(diào)研發(fā)現(xiàn)，不僅百度的paddle框架支持混合精度訓練，在Tensorflow和Pytorch中也有相應的實現(xiàn)。下面我們先來講講理論，后面再分析混合精度訓練在三大深度學習框架中的打開方式。

理論原理

訓練過神經(jīng)網(wǎng)絡的小伙伴都知道，神經(jīng)網(wǎng)絡的參數(shù)和中間結(jié)果絕大部分都是單精度浮點數(shù)（即float32）存儲和計算的，當網(wǎng)絡變得超級大時，降低浮點數(shù)精度，比如使用半精度浮點數(shù)，顯然是提高計算速度，降低存儲開銷的一個很直接的辦法。

然而副作用也很顯然，如果我們直接降低浮點數(shù)的精度直觀上必然導致模型訓練精度的損失。但是呢，天外有天，這篇文章用了三種機制有效地防止了模型的精度損失。待小夕一一說來o(*￣▽￣*)ブ

權(quán)重備份(master weights)

我們知道半精度浮點數(shù)（float16）在計算機中的表示分為1bit的符號位，5bits的指數(shù)位和10bits的尾數(shù)位，所以它能表示的最小的正數(shù)即2^-24（也就是精度到此為止了）。當神經(jīng)網(wǎng)絡中的梯度灰常小的時候，網(wǎng)絡訓練過程中每一步的迭代（灰常小的梯度 ✖ 也黑小的learning rate）會變得更小，小到float16精度無法表示的時候，相應的梯度就無法得到更新。

論文統(tǒng)計了一下在Mandarin數(shù)據(jù)集上訓練DeepSpeech 2模型時產(chǎn)生過的梯度，發(fā)現(xiàn)在未乘以learning rate之前，就有接近5%的梯度直接悲劇的變成0（精度比2^-24還要高的梯度會直接變成0），造成重大的損失呀/(ㄒoㄒ)/~~

還有更難的，假設迭代量逃過一劫準備奉獻自己的時候。。。由于網(wǎng)絡中的權(quán)重往往遠大于我們要更新的量，當?shù)啃∮贔loat16當前區(qū)間內(nèi)能表示的最小間隔的時候，更新也會失敗（哭瞎┭┮﹏┭┮我怎么這么難鴨）

所以怎么辦呢？作者這里提出了一個非常simple but effective的方法，就是前向傳播和梯度計算都用float16，但是存儲網(wǎng)絡參數(shù)的梯度時要用float32！這樣就可以一定程度上的解決上面說的兩個問題啦~~~

我們來看一下訓練曲線，藍色的線是正常的float32精度訓練曲線，橙色的線是使用float32存儲網(wǎng)絡參數(shù)的learning curve，綠色滴是不使用float32存儲參數(shù)的曲線，兩者一比就相形見絀啦。

損失放縮（loss scaling）

有了上面的master weights已經(jīng)可以足夠高精度的訓練很多網(wǎng)絡啦，但是有點強迫癥的小夕來說怎么還是覺得有點不對呀o((⊙﹏⊙))o.

雖然使用float32來存儲梯度，確實不會丟失精度了，但是計算過程中出現(xiàn)的指數(shù)位小于 -24 的梯度不還是會丟失的嘛！相當于用漏水的篩子從河邊往村里運水，為了多存點水，村民們把儲水的碗換成了大缸，燃鵝篩子依然是漏的哇，在路上的時候水就已經(jīng)漏的木有了。。

于是loss scaling方法來了。首先作者統(tǒng)計了一下訓練過程中激活函數(shù)梯度的分布情況，由于網(wǎng)絡中的梯度往往都非常小，導致在使用FP16的時候右邊有大量的范圍是沒有使用的。這種情況下， 我們可以通過放大loss來把整個梯度右移，減少因為精度隨時變?yōu)?的梯度。

那么問題來了，怎么合理的放大loss呢？一個最簡單的方法是常數(shù)縮放，把loss一股腦統(tǒng)一放大S倍。float16能表示的最大正數(shù)是2^15*(1+1-2^-10)=65504，我們可以統(tǒng)計網(wǎng)絡中的梯度，計算出一個常數(shù)S，使得最大的梯度不超過float16能表示的最大整數(shù)即可。

當然啦，還有更加智能的動態(tài)調(diào)整(automatic scaling) o(*￣▽￣*)ブ

我們先初始化一個很大的S，如果梯度溢出，我們就把S縮小為原來的二分之一；如果在很多次迭代中梯度都沒有溢出，我們也可以嘗試把S放大兩倍。以此類推，實現(xiàn)動態(tài)的loss scaling。

運算精度（precison of ops）

精益求精再進一步，神經(jīng)網(wǎng)絡中的運算主要可以分為四大類，混合精度訓練把一些有更高精度要求的運算，在計算過程中使用float32，存儲的時候再轉(zhuǎn)換為float16。

• matrix multiplication: linear, matmul, bmm, conv
• pointwise: relu, sigmoid, tanh, exp, log
• reductions: batch norm, layer norm, sum, softmax
• loss functions: cross entropy, l2 loss, weight decay

像矩陣乘法和絕大多數(shù)pointwise的計算可以直接使用float16來計算并存儲，而reductions、loss function和一些pointwise（如exp，log，pow等函數(shù)值遠大于變量的函數(shù)）需要更加精細的處理，所以在計算中使用用float32，再將結(jié)果轉(zhuǎn)換為float16來存儲。

總結(jié)：三大深度學習框架的打開方式

混合精度訓練做到了在前向和后向計算過程中均使用半精度浮點數(shù)，并且沒有像之前的一些工作一樣還引入額外超參，而且重要的是，實現(xiàn)非常簡單卻能帶來非常顯著的收益，在顯存half以及速度double的情況下保持模型的精度，簡直不能再厲害啦。

看完了硬核技術(shù)細節(jié)之后，我們趕緊來看看代碼實現(xiàn)吧！如此強大的混合精度訓練的代碼實現(xiàn)不要太簡單了吧😮

Pytorch

導入Automatic Mixed Precision (AMP)，不要998不要288，只需3行無痛使用！

from apex import ampmodel, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 這里是“歐一”，不是“零一”with amp.scale_loss(loss, optimizer) as scaled_loss:scaled_loss.backward() 

來看個例子，將上面三行按照正確的位置插入到自己原來的代碼中就可以實現(xiàn)酷炫的半精度訓練啦！

import torchfrom apex import ampmodel = ... optimizer = ...#包裝model和optimizermodel, optimizer = amp.initialize(model, optimizer, opt_level="O1")for data, label in data_iter: out = model(data) loss = criterion(out, label) optimizer.zero_grad() #loss scaling，代替loss.backward() with amp.scaled_loss(loss, optimizer) as scaled_loss:scaled_loss.backward() optimizer.step() 

Tensorflow

一句話實現(xiàn)混合精度訓練之修改環(huán)境變量，在python腳本中設置環(huán)境變量

os.environ[ TF_ENABLE_AUTO_MIXED_PRECISION ] = 1 

除此之外，也可以用類似pytorch的方式來包裝optimizer。

Graph-based示例

opt = tf.train.AdamOptimizer()#add a lineopt = tf.train.experimental.enable_mixed_precision_graph_rewrite( opt, loss_scale= dynamic ) train_op = opt.miminize(loss) 

Keras-based示例

opt = tf.keras.optimizers.Adam()#add a lineopt = tf.train.experimental.enable_mixed_precision_graph_rewrite( opt, loss_scale= dynamic ) model.compile(loss=loss, optimizer=opt)model.fit(...) 

PaddlePaddle

一句話實現(xiàn)混合精度訓練之添加config（驚呆🙃畢竟混合精度訓練是百度家提出的，內(nèi)部早就熟練應用了叭）

--use_fp16=true 

舉個栗子，基于BERT finetune XNLI任務時，只需在執(zhí)行時設置use_fp16為true即可。

export FLAGS_sync_nccl_allreduce=0export FLAGS_eager_delete_tensor_gb=1export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7BERT_BASE_PATH="chinese_L-12_H-768_A-12"TASK_NAME= XNLI DATA_PATH=/path/to/xnli/data/CKPT_PATH=/path/to/save/checkpoints/python -u run_classifier.py --task_name ${TASK_NAME} --use_fp16=true #!!!!!!add a line --use_cuda true --do_train true --do_val true --do_test true --batch_size 32 --in_tokens false --init_pretraining_params ${BERT_BASE_PATH}/params --data_dir ${DATA_PATH} --vocab_path ${BERT_BASE_PATH}/vocab.txt --checkpoints ${CKPT_PATH} --save_steps 1000 --weight_decay 0.01 --warmup_proportion 0.1 --validation_steps 100 --epoch 3 --max_seq_len 128 --bert_config_path ${BERT_BASE_PATH}/bert_config.json --learning_rate 5e-5 --skip_steps 10 --num_iteration_per_drop_scope 10 --verbose true