[Pytorch]基于混和精度的模型加速
这篇博客是在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到很多更加具体地内容。
感谢团队同学推荐。
