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

 

以上