Pytorch 1.6使用自动混合精度训练（AMP）

今天pytorch 1.6正式发布了，本次更新的亮点在于引入了自动混合精度训练，详情见官网https://pytorch.org/blog/pytorch-1.6-released/

在此做一下简介

自动混合精度的意义在于加入了半精度的张量类型，这种类型可以在某些运算中具有更快的速度（如卷积和全连接层），官方文档中支持半精度的类型如下

__matmul__, addbmm, addmm, addmv, addr, baddbmm, bmm, chain_matmul, 

conv1d, conv2d, conv3d, conv_transpose1d, conv_transpose2d, conv_transpose3d, 

linear, matmul, mm, mv, prelu

此部分详细介绍见https://pytorch.org/docs/stable/amp.html#autocast-op-reference

有些暂时不支持半精度，比如LSTM和GRU，使用会报错，提示

RuntimeError: cuDNN error: CUDNN_STATUS_BAD_PARAM

 

【使用方法】

一个标准例子如下

from torch.cuda.amp import autocast as autocast, GradScaler

# 创建model，默认是torch.FloatTensor
model = Net().cuda()
optimizer = optim.SGD(model.parameters(), ...)

# 在训练最开始之前实例化一个GradScaler对象
scaler = GradScaler()

for epoch in epochs:
    for input, target in data:
        optimizer.zero_grad()

        # 前向过程(model + loss)开启 autocast
        with autocast():
            output = model(input)
            loss = loss_fn(output, target)

        # Scales loss，这是因为半精度的数值范围有限，因此需要用它放大
        scaler.scale(loss).backward()

        # scaler.step() unscale之前放大后的梯度，但是scale太多可能出现inf或NaN
        # 故其会判断是否出现了inf/NaN
        # 如果梯度的值不是 infs 或者 NaNs, 那么调用optimizer.step()来更新权重,
        # 如果检测到出现了inf或者NaN，就跳过这次梯度更新，同时动态调整scaler的大小
        scaler.step(optimizer)

        # 查看是否要更新scaler
        scaler.update()

更多实例以及Doc见官网

https://pytorch.org/docs/stable/notes/amp_examples.html#typical-mixed-precision-training

 

关于半精度的详细解释，推荐一篇知乎文章，见

PyTorch的自动混合精度（AMP） - Gemfield的文章 - 知乎 https://zhuanlan.zhihu.com/p/165152789

 

以上