Pytorch >= 1.6.0

pytorch1.6开始官方支持了混合精度训练。使用方法如下：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
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()

Pytorch < 1.6.0

pytorch1.6版本前需要配合第三方插件使用混合精度。

Apex

Apex是由Nvidia维护的一个支持混合精度分布式训练的第三方Pytorch扩展。可以用短短三行代码就能实现不同程度的混合精度加速，使训练时间和显存占用直接缩小一半。

安装

以下是安装 apex 的方法. apex克隆在哪里都无所谓

1
2
3
git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--pyprof" --global-option="--cpp_ext" --global-option="--cuda_ext" ./

以下代码用于验证是否安装成功：

1
2
3
4
5
6
7
try:
    from apex.parallel import DistributedDataParallel as DDP
    from apex.fp16_utils import *
    from apex import amp, optimizers
    from apex.multi_tensor_apply import multi_tensor_applier
except ImportError:
    raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")

卸载命令为：

1
pip uninstall apex

支持列表

判断你的GPU是否支持FP16：支持的有拥有Tensor Core的GPU(2080Ti、Titan、Tesla等)，不支持的(Pascal系列)就不建议折腾了。

如果你的GPU是以下GPU的其中一种: 请调整nvcc与pytorch.cuda至10.0

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
GeForce GTX 1650
GeForce GTX 1660
GeForce GTX 1660 Ti
GeForce RTX 2060
GeForce RTX 2060 Super
GeForce RTX 2070
GeForce RTX 2070 Super
GeForce RTX 2080
GeForce RTX 2080 Super
GeForce RTX 2080 Ti
Titan RTX
Quadro RTX 4000
Quadro RTX 5000
Quadro RTX 6000
Quadro RTX 8000
Tesla T4

如果你打算调整Pytorch Version来适应APEX(推荐)

首先用 nvcc --version 命令调查你的nvcc版本

1
2
3
4
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2018 NVIDIA Corporation
Built on Tue_Jun_12_23:07:04_CDT_2018
Cuda compilation tools, release 9.2, V9.2.148

再进入python看看pytorch cuda版本

1
2
>>> torch.version.cuda
'9.0.176'

问题解决

3090ti自带cuda为11.1，但是并没有匹配的torch版本，这里为了安上apex，可以通过修改setup.py尝试禁用次要版本检查。

找到如下函数，注释掉原语句，返回空值：

1
2
 def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
+    return

修改后，重新安装apex即可：

1
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

我在安装了带有cuda-11.0的pytorch的同时，针对整个系统的cuda-11.1成功构建了apex！

使用

导入包：

1
from apex import amp

在你创建后net(model)和optimizer后添加代码(注意顺序)：

1
2
3
4
5
net = net.cuda()
net, optimizer = amp.initialize(net, optimizer, opt_level="O1") # 这里是“欧一”，不是“零一”

# if you need to use multiple GPU, uncomment this line
# net = torch.nn.DataParallel(net, device_ids=[i for i in range(torch.cuda.device_count())])

然后把你的loss.backward()换成 ：

1
2
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

如果你有用Pytorch自带的BCELoss或F.binary_cross_entropy, 请把他们替换成BCEWithLogitsLoss或F.binary_cross_entropy_with_logists并用logists 作为输入.

说明

opt level

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

  • oe：纯FP32训练，可以作为accuracy的baseline；

  • o1：混合精度训练(推荐使用)，根据黑白名单自动决定使用FP16(GEMM，卷积)还是FP32(Softmax)进行计算。

  • 02：“几乎FP16”混合精度训练，不存在黑白名单，除了Batch norm，几乎都是用FP16计算。

  • o3：纯FP16训练，很不稳定，但是可以作为speed的baseline；

动态损失放大（Dynamic Loss Scaling）

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

优势

  1. 减少显存占用 现在模型越来越大，当你使用Bert这一类的预训练模型时，往往显存就被模型及模型计算占去大半，当想要使用更大的Batch Size的时候会显得捉襟见肘。由于FP16的内存占用只有FP32的一半，自然地就可以帮助训练过程节省一半的显存空间。
  2. 加快训练和推断的计算 与普通的空间时间Trade-off的加速方法不同，FP16除了能节约内存，还能同时节省模型的训练时间。在大部分的测试中，基于FP16的加速方法能够给模型训练带来多一倍的加速体验（爽感类似于两倍速看肥皂剧）。
  3. 张量核心的普及 硬件的发展同样也推动着模型计算的加速，随着Nvidia张量核心(Tensor Core)的普及，16bit计算也一步步走向成熟，低精度计算也是未来深度学习的一个重要趋势，再不学习就out啦。