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
    return

4、自动进行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