3行代碼提速模型訓練：這個算法讓你的GPU老樹開新花
百度和Nvidia研究院結合N卡底層計算優化，提出了一種有效的神經網絡訓練加速方法，不僅是預訓練，在全民finetune BERT的今天變得異常有用。
一切還要從2018年ICLR的一篇論文說起。
《MIXED PRECISION TRAINING》是百度&Nvidia研究院一起發表的，結合N卡底層計算優化，提出了一種灰常有效的神經網絡訓練加速方法，不僅是預訓練，在全民finetune BERT的今天變得異常有用哇。
而且調研發現，不僅百度的paddle框架支持混合精度訓練，在Tensorflow和Pytorch中也有相應的實現。下面我們先來講講理論，后面再分析混合精度訓練在三大深度學習框架中的打開方式。
理論原理
訓練過神經網絡的小伙伴都知道，神經網絡的參數和中間結果絕大部分都是單精度浮點數（即float32）存儲和計算的，當網絡變得超級大時，降低浮點數精度，比如使用半精度浮點數，顯然是提高計算速度，降低存儲開銷的一個很直接的辦法。
然而副作用也很顯然，如果我們直接降低浮點數的精度直觀上必然導致模型訓練精度的損失。但是呢，天外有天，這篇文章用了三種機制有效地防止了模型的精度損失。待小夕一一說來o(*￣▽￣*)ブ
權重備份(master weights)
我們知道半精度浮點數（float16）在計算機中的表示分為1bit的符號位，5bits的指數位和10bits的尾數位，所以它能表示的最小的正數即2^-24（也就是精度到此為止了）。當神經網絡中的梯度灰常小的時候，網絡訓練過程中每一步的迭代（灰常小的梯度 ✖ 也黑小的learning rate）會變得更小，小到float16精度無法表示的時候，相應的梯度就無法得到更新。
論文統計了一下在Mandarin數據集上訓練DeepSpeech 2模型時產生過的梯度，發現在未乘以learning rate之前，就有接近5%的梯度直接悲劇的變成0（精度比2^-24還要高的梯度會直接變成0），造成重大的損失呀/(ㄒoㄒ)/~~
還有更難的，假設迭代量逃過一劫準備奉獻自己的時候。。。由于網絡中的權重往往遠大于我們要更新的量，當迭代量小于Float16當前區間內能表示的最小間隔的時候，更新也會失?。尴供雪药n┭┮我怎么這么難鴨）
所以怎么辦呢？作者這里提出了一個非常simple but effective的方法，就是前向傳播和梯度計算都用float16，但是存儲網絡參數的梯度時要用float32！這樣就可以一定程度上的解決上面說的兩個問題啦~~~
我們來看一下訓練曲線，藍色的線是正常的float32精度訓練曲線，橙色的線是使用float32存儲網絡參數的learning curve，綠色滴是不使用float32存儲參數的曲線，兩者一比就相形見絀啦。
損失放縮（loss scaling）
有了上面的master weights已經可以足夠高精度的訓練很多網絡啦，但是有點強迫癥的小夕來說怎么還是覺得有點不對呀o((⊙﹏⊙))o.
雖然使用float32來存儲梯度，確實不會丟失精度了，但是計算過程中出現的指數位小于 -24 的梯度不還是會丟失的嘛！相當于用漏水的篩子從河邊往村里運水，為了多存點水，村民們把儲水的碗換成了大缸，燃鵝篩子依然是漏的哇，在路上的時候水就已經漏的木有了。。
于是loss scaling方法來了。首先作者統計了一下訓練過程中激活函數梯度的分布情況，由于網絡中的梯度往往都非常小，導致在使用FP16的時候右邊有大量的范圍是沒有使用的。這種情況下， 我們可以通過放大loss來把整個梯度右移，減少因為精度隨時變為0的梯度。
那么問題來了，怎么合理的放大loss呢？一個最簡單的方法是常數縮放，把loss一股腦統一放大S倍。float16能表示的最大正數是2^15*(1+1-2^-10)=65504，我們可以統計網絡中的梯度，計算出一個常數S，使得最大的梯度不超過float16能表示的最大整數即可。
當然啦，還有更加智能的動態調整(automatic scaling) o(*￣▽￣*)ブ
我們先初始化一個很大的S，如果梯度溢出，我們就把S縮小為原來的二分之一；如果在很多次迭代中梯度都沒有溢出，我們也可以嘗試把S放大兩倍。以此類推，實現動態的loss scaling。
運算精度（precison of ops）
精益求精再進一步，神經網絡中的運算主要可以分為四大類，混合精度訓練把一些有更高精度要求的運算，在計算過程中使用float32，存儲的時候再轉換為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
像矩陣乘法和絕大多數pointwise的計算可以直接使用float16來計算并存儲，而reductions、loss function和一些pointwise（如exp，log，pow等函數值遠大于變量的函數）需要更加精細的處理，所以在計算中使用用float32，再將結果轉換為float16來存儲。
總結：三大深度學習框架的打開方式
混合精度訓練做到了在前向和后向計算過程中均使用半精度浮點數，并且沒有像之前的一些工作一樣還引入額外超參，而且重要的是，實現非常簡單卻能帶來非常顯著的收益，在顯存half以及速度double的情況下保持模型的精度，簡直不能再厲害啦。
看完了硬核技術細節之后，我們趕緊來看看代碼實現吧！如此強大的混合精度訓練的代碼實現不要太簡單了吧😮
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()
來看個例子，將上面三行按照正確的位置插入到自己原來的代碼中就可以實現酷炫的半精度訓練啦！
- 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
一句話實現混合精度訓練之修改環境變量，在python腳本中設置環境變量
- 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
一句話實現混合精度訓練之添加config（驚呆🙃畢竟混合精度訓練是百度家提出的，內部早就熟練應用了叭）
- --use_fp16=true
舉個栗子，基于BERT finetune XNLI任務時，只需在執行時設置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
