混合精度训练

原文来自于机器学习算法与自然语言处理公众号

混合精度训练

• 《MIXED PRECISION TRAINING》这篇论文是百度&Nvidia研究院一起发表的，结合N卡底层计算优化，提出了一种灰常有效的神经网络训练加速方法，不仅是预训练，在全民finetune BERT的今天变得异常有用。而不仅百度的paddle框架支持混合精度训练，在Tensorflow和Pytorch中也有相应的实现。下面我们先来讲讲理论，后面再分析混合精度训练在三大深度学习框架中的打开方式。

理论原理

• 训练过神经网络的小伙伴都知道，神经网络的参数和中间结果绝大部分都是单精度浮点数（即float32）存储和计算的，当网络变得超级大时，降低浮点数精度，比如使用半精度浮点数，显然是提高计算速度，降低存储开销的一个很直接的办法。然而副作用也很显然，如果我们直接降低浮点数的精度直观上必然导致模型训练精度的损失。但是呢，天外有天，这篇文章用了三种机制有效地防止了模型的精度损失。

权重备份(master weights)

• 我们知道半精度浮点数（float16）在计算机中的表示分为1bit的符号位，5bits的指数位和10bits的尾数位，所以它能表示的最小的正数即2^-24（也就是精度到此为止了）。当神经网络中的梯度灰常小的时候，网络训练过程中每一步的迭代（灰常小的梯度 ✖ 也黑小的learning rate）会变得更小，小到float16精度无法表示的时候，相应的梯度就无法得到更新。

• 论文统计了一下在Mandarin数据集上训练DeepSpeech 2模型时产生过的梯度，发现在未乘以learning rate之前，就有接近5%的梯度直接悲剧的变成0（精度比2^-24还要高的梯度会直接变成0），造成重大的损失。

  
  

• 还有更难的，假设迭代量逃过一劫准备奉献自己的时候。。。由于网络中的权重往往远大于我们要更新的量，当迭代量小于Float16当前区间内能表示的最小间隔的时候，更新也会失败。
• 作者这里提出了一个非常simple but effective的方法，就是前向传播和梯度计算都用float16，但是存储网络参数的梯度时要用float32！这样就可以一定程度上的解决上面说的两个问题啦。

• 我们来看一下训练曲线，蓝色的线是正常的float32精度训练曲线，橙色的线是使用float32存储网络参数的learning curve，绿色滴是不使用float32存储参数的曲线，两者一比就相形见绌啦。

损失放缩（loss scaling）

• 虽然使用float32来存储梯度，确实不会丢失精度了，但是计算过程中出现的指数位小于 -24 的梯度不还是会丢失。于是loss scaling方法来了。首先作者统计了一下训练过程中激活函数梯度的分布情况，由于网络中的梯度往往都非常小，导致在使用FP16的时候右边有大量的范围是没有使用的。这种情况下， 我们可以通过放大loss来把整个梯度右移，减少因为精度随时变为0的梯度。
• 那么问题来了，怎么合理的放大loss呢？一个最简单的方法是常数缩放，把loss一股脑统一放大S倍。float16能表示的最大正数是2^15*(1+1-2-10)=65504，我们可以统计网络中的梯度，计算出一个常数S，使得最大的梯度不超过float16能表示的最大整数即可。
• 当然啦，还有更加智能的动态调整(automatic scaling) 我们先初始化一个很大的S，如果梯度溢出，我们就把S缩小为原来的二分之一；如果在很多次迭代中梯度都没有溢出，我们也可以尝试把S放大两倍。以此类推，实现动态的loss scaling。

运算精度（precison of ops）

• 精益求精再进一步，神经网络中的运算主要可以分为四大类，混合精度训练把一些有更高精度要求的运算，在计算过程中使用float32，存储的时候再转换为float16。
• 像矩阵乘法和绝大多数pointwise的计算可以直接使用float16来计算并存储，而reductions、loss function和一些pointwise（如exp，log，pow等函数值远大于变量的函数）需要更加精细的处理，所以在计算中使用用float32，再将结果转换为float16来存储。

Pytorch

• 导入Automatic Mixed Precision (AMP)

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()

import torch
from apex import amp
model = ... 
optimizer = ...

#包装model和optimizer
model, 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'

• Graph-based示例

opt = tf.train.AdamOptimizer()

#add a line
opt = 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 line
opt = 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=0
export FLAGS_eager_delete_tensor_gb=1
export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7

BERT_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