這篇博客是在pytorch中基於apex使用混合精度加速的一個偏工程的描述，原理層面的解釋並不是這篇博客的目的，不過在參考部分提供了非常有價值的資料，可以進一步研究。
一個關鍵原則：“僅僅在權重更新的時候使用fp32，耗時的前向和后向運算都使用fp16”。其中的一個技巧是：在反向計算開始前，將dloss乘上一個scale，人為變大；權重更新前，除去scale，恢復正常值。目的是為了減小激活gradient下溢出的風險。
apex是nvidia的一個pytorch擴展，用於支持混合精度訓練和分布式訓練。在之前的博客中，神經網絡的Low-Memory技術梳理了一些low-memory技術，其中提到半精度，比如fp16。apex中混合精度訓練可以通過簡單的方式開啟自動化實現，組里同學交流的結果是：一般情況下，自動混合精度訓練的效果不如手動修改。分布式訓練中，有社區同學心心念念的syncbn的支持。關於syncbn，在去年做CV的時候，我們就有一些來自民間的嘗試，不過具體提升還是要考慮具體任務場景。
那么問題來了，如何在pytorch中使用fp16混合精度訓練呢？
第零：混合精度訓練相關的參數
parser.add_argument('--fp16',
action='store_true',
help="Whether to use 16-bit float precision instead of 32-bit")
parser.add_argument('--loss_scale',
type=float, default=0,
help="Loss scaling to improve fp16 numeric stability. Only used when fp16 set to True.\n"
"0 (default value): dynamic loss scaling.\n"
"Positive power of 2: static loss scaling value.\n")
第一：模型參數轉換為fp16
nn.Module中的half()方法將模型中的float32轉化為float16，實現的原理是遍歷所有tensor，而float32和float16都是tensor的屬性。也就是說，一行代碼解決，如下：
model.half()
第二：修改優化器
在pytorch下，當使用fp16時，需要修改optimizer。類似代碼如下（代碼參考這里）：
# Prepare optimizer
if args.do_train:
param_optimizer = list(model.named_parameters())
no_decay = ['bias', 'LayerNorm.bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [
{'params': [p for n, p in param_optimizer if not any(nd in n for nd in no_decay)], 'weight_decay': 0.01},
{'params': [p for n, p in param_optimizer if any(nd in n for nd in no_decay)], 'weight_decay': 0.0}
]
if args.fp16:
try:
from apex.optimizers import FP16_Optimizer
from apex.optimizers import FusedAdam
except ImportError:
raise ImportError("Please install apex from https://www.github.com/nvidia/apex to use distributed and fp16 training.")
optimizer = FusedAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
bias_correction=False,
max_grad_norm=1.0)
if args.loss_scale == 0:
optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
else:
optimizer = FP16_Optimizer(optimizer, static_loss_scale=args.loss_scale)
warmup_linear = WarmupLinearSchedule(warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
else:
optimizer = BertAdam(optimizer_grouped_parameters,
lr=args.learning_rate,
warmup=args.warmup_proportion,
t_total=num_train_optimization_steps)
第三：backward時做對應修改
if args.fp16:
optimizer.backward(loss)
else:
loss.backward()
第四：學習率修改
if args.fp16:
# modify learning rate with special warm up BERT uses
# if args.fp16 is False, BertAdam is used that handles this automatically
lr_this_step = args.learning_rate * warmup_linear.get_lr(global_step, args.warmup_proportion)
for param_group in optimizer.param_groups:
param_group['lr'] = lr_this_step
optimizer.step()
optimizer.zero_grad()
根據參考3，值得重述一些重要結論：
（1）深度學習訓練使用16bit表示/運算正逐漸成為主流。
（2）低精度帶來了性能、功耗優勢，但需要解決量化誤差（溢出、舍入）。
（3）常見的避免量化誤差的方法：為權重保持高精度（fp32)備份；損失放大，避免梯度的下溢出；一些特殊層（如BatchNorm）仍使用fp32運算。
參考資料：
1.nv官方repo給了一些基於pytorch的apex加速的實現
實現是基於fairseq實現的，可以直接對比代碼1-apex版和代碼2-非apex版(fairseq官方版)，了解是如何基於apex實現加速的。
按圖索驥，可以get到很多更加具體地內容。
感謝團隊同學推薦。