pytorch 混合精度训练
半精度混合训练，省显存且能提速一倍。混合精度训练，需要硬件设备支持才可以。混合精度需要Tensor Core支持，P4卡不支持混合精度训练。
Tensor Core：我们知道在深度学习中大量的运算都是在高维矩阵（张量 Tensor）间进行的，Tensor Core可以让tensor的计算速度急速上升。Tensor Core专业执行矩阵数学运算，适用于深度学习和某些类型的HPC。
Tensor Core执行融合乘法加法，其中两个44 FP16矩阵相乘，然后将结果添加到44 FP16或FP32矩阵中，最终输出新的4*4 FP16或FP32矩阵。NVIDIA将Tensor Core进行的这种运算称为混合精度数学，因为输入矩阵的精度为半精度，但乘积可以达到完全精度。
使用半精度混合训练满足的条件：
判断你的GPU是否支持FP16：构拥有Tensor Core的GPU（2080Ti、Titan、Tesla等），不支持的（Pascal系列）就不建议折腾了。
常数的范围：为了保证计算不溢出，首先要保证人为设定的常数（包括调用的源码中的）不溢出，如各种epsilon，INF等。
Dimension最好是8的倍数：Nvidia官方的文档的2.2条表示，维度都是8的倍数的时候，性能最好。
涉及到sum的操作要小心，很容易溢出，类似Softmax的操作建议用官方API，并定义成layer写在模型初始化里。
模型书写要规范：自定义的Layer写在模型初始化函数里，graph计算写在forward里。
某些不常用的函数，在使用前需要注册：amp.register_float_function(torch, 'sigmoid')
某些函数（如einsum）暂不支持FP16加速，建议不要用的太heavy，xlnet的实现改FP16困扰了我很久。
需要操作模型参数的模块（类似EMA），要使用AMP封装后的model。
需要操作梯度的模块必须在optimizer的step里，不然AMP不能判断grad是否为Nan。
常见的支持混合精度显卡：
V100、GTX 2080Ti、Titan。
实际中使用autocast + GradScaler实现混合精度训练。
1、autocast
from torch.cuda.amp import autocast as autocast
# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
for input, target in data:
optimizer.zero_grad()
# 前向过程(model + loss)开启 autocast
with autocast():
output = model(input)
loss = loss_fn(output, target)
# 反向传播在autocast上下文之外
loss.backward()
optimizer.step()2、GradScaler
GradScaler就是梯度scaler模块，需要在训练最开始之前实例化一个GradScaler对象。来自HRNet工程代码。
from torch.cuda.amp import autocast as autocast
# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)
# 在训练最开始之前实例化一个GradScaler对象
scaler = torch.cuda.amp.GradScaler() if args.amp else None
for epoch in epochs:
for input, target in data:
optimizer.zero_grad()
# 前向过程(model + loss)开启 autocast
with torch.cuda.amp.autocast(enabled=scaler is not None):
results = model(images)
losses = mse(results, targets)
optimizer.zero_grad()
if scaler is not None:
scaler.scale(losses).backward()
scaler.step(optimizer)
scaler.update()
else:
losses.backward()
optimizer.step()3、nn.DataParallel
class Model(nn.Module):
def __init__(self):
super(Model, self).__init__()
def forward(self, input_data_c1):
with autocast():
# code
return4、自动进行autocast的操作
如下操作中tensor会被自动转化为半精度浮点型的torch.HalfTensor：
matmul
addbmm
addmm
addmv
addr
baddbmm
bmm
chain_matmul
conv1d
conv2d
conv3d
conv_transpose1d
conv_transpose2d
conv_transpose3d
linear
matmul
mm
mv
prelu其他操作比如rnn也可以进行半精度运行，但是需要自己手动，暂时没有提供自动的转换。
阅读：
1、混合精度训练
极市开发者平台-计算机视觉算法开发落地平台
2、深度学习中GPU卡的选择
深度学习中GPU卡的选择 - 知乎
3、混合精度训练-实验
混合精度训练-实验 - 知乎
4、英伟达显卡有这么多型号，怎么选择？（选显卡提供依据参考）
英伟达显卡有这么多型号，运行深度学习哪一个型号性价比最高？是算力越大越好么？ - 知乎
5、PyTorch 半精度训练踩坑
https://www.jianshu.com/p/8d7831d01cf4