Pytorch混合精度(FP16&FP32)(AMP自动混合精度)/半精度 训练(二) —— 代码示例 apex pytorch

      现在已经有内置的torch.cuda.amp了，可以不用apex了

      为了帮助提高Pytorch的训练效率，NVIDIA开源了一款混合精度训练工具Apex。号称能够在不降低性能的情况下，将模型训练的速度提升2-4倍，训练显存消耗减少为之前的一半。

      该工具提供了三个功能，amp、parallel和normalization  

      Apex能够降低训练精度提高batch大小，512*512的图片在8G的显存上batch只能为1，2，使用了apex后可以增大到10。只支持Pytorch。

代码实现

简单来说，就是三句话

from apex import amp
model, optimizer = amp.initialize(model, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

1. opt_level

其中只有一个opt_level需要用户自行配置：

• O0：纯FP32训练，可以作为accuracy的baseline；
• O1：混合精度训练（推荐使用），根据黑白名单自动决定使用FP16（GEMM, 卷积）还是FP32（Softmax）进行计算。
• O2：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算。
• O3：纯FP16训练，很不稳定，但是可以作为speed的baseline；

2. 动态损失放大（Dynamic Loss Scaling）

      AMP默认使用动态损失放大。为了充分利用FP16的范围，缓解舍入误差，尽量使用最高的放大倍数（2^24），如果产生了上溢出（Overflow），则跳过参数更新，缩小放大倍数使其不溢出，在一定步数后（比如2000步）会再尝试使用大的scale来充分利用FP16的范围

普通代码

import torch
N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device="cuda")
y = torch.randn(N, D_out, device="cuda")
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
for to in range(500):
   y_pred = model(x)
   loss = torch.nn.functional.mse_loss(y_pred, y)
   optimizer.zero_grad()
   loss.backward()
   optimizer.step()

添加amp之后

import torch
from apex import amp

N, D_in, D_out = 64, 1024, 512
x = torch.randn(N, D_in, device="cuda")
y = torch.randn(N, D_out, device="cuda")
model = torch.nn.Linear(D_in, D_out).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
for to in range(500):
    y_pred = model(x)
    loss = torch.nn.functional.mse_loss(y_pred, y)
    optimizer.zero_grad()
    with amp.scale_loss(loss, optimizer) as scaled_loss:
        scaled_loss.backward()
    optimizer.step()   

而且在使用的时候，需要判断你的GPU是否支持FP16：支持的有拥有Tensor Core的GPU（2080Ti、Titan、Tesla等），不支持的（Pascal系列）就不建议折腾了。

跑amp要注意的坑

知乎上有人说，“现在在做深度学习模型的时候，几乎都会第一时间把代码改成混合精度训练的了，速度快，精度还不减，确实是调参炼丹必备神器”。