使用AMP的PyTorch模型更快，内存效率更高

作者 | Rahul Agarwal

来源 | Medium

编辑 | 代码医生团队

您是否知道反向传播算法是Geoffrey Hinton 在1986年的《自然》杂志上提出的？

同样，卷积网络由Yann le cun于1998年首次提出进行数字分类，他使用了单个卷积层。直到2012年下半年，Alexnet才通过使用多个卷积层在imagenet上实现最先进的技术来推广Convnets。

那么，是什么让他们现在而不是以前如此著名？

只有在拥有大量计算资源的情况下，才能实验并充分利用深度学习在最近的全部潜力。

但是是否充分利用了计算资源？可以做得更好吗？

这篇文章是关于利用Tensor核心和自动混合精度来更快地训练深度学习网络的。

什么是张量核心？

根据NVIDIA网站：

NVIDIA Turing和Volta GPU由Tensor Cores提供支持，Tensor Cores是一项革命性技术，可提供开创性的AI性能。Tensor Core可以加速AI核心的大型矩阵运算，并在单个运算中执行混合精度矩阵乘法和累加计算。在一个NVIDIA GPU中并行运行数百个Tensor Core，这可以极大地提高吞吐量和效率

简单地说; 它们是专门用于特定类型矩阵运算的专用内核。

可以将两个FP16矩阵相乘并将其添加到FP16 / FP32矩阵中，从而得到FP16 / FP32矩阵。Tensor内核支持混合精度数学，即输入为半精度（FP16），输出为全精度（FP32）。上面的操作对于许多深度学习任务具有内在的价值，并且Tensor内核为该操作提供了专用的硬件。

现在，使用FP16和FP32主要有两个好处。

1. FP16需要较少的内存，因此更易于训练和部署大型神经网络。它还减少了数据移动。
2. 使用Tensor Core，数学运算的运行速度大大降低了精度。NVIDIA提供的Volta GPU的确切数量是：FP16中为125 TFlops，而FP32中为15.7 TFlops（加速8倍）

但是也有缺点。从FP32转到FP16时，必然会降低精度。

FP32与FP16：FP32具有八个指数位和23个小数位，而FP16具有五个指数位和十个小数位。

但是需要FP32吗？

FP16实际上可以很好地表示大多数权重和渐变。因此，拥有存储和使用FP32所需的所有这些额外位只是浪费。

那么如何使用Tensor Core？

检查了Titan RTX GPU是否具有576张量内核以及4,608个NVIDIA CUDA内核。但是如何使用这些张量核心？

坦白地说，NVIDIA可以轻松地将Tensor内核与自动混合精度一起使用，并提供了几行代码。需要在代码中做两件事：

1.FP32所需的操作（如Softmax）被分配给FP32，而FP16可以完成的操作（如Conv）被自动分配给FP16。

2.使用损耗定标保留较小的梯度值。梯度值可能超出FP16的范围。在这种情况下，将对梯度值进行缩放，使其落在FP16范围内。

如果还不了解背景细节，那也可以。代码实现相对简单。

使用PyTorch进行混合精度训练：

从PyTorch中的基本网络开始。

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()

要利用自动混合精度训练的优势，首先需要安装apex库。只需在终端中运行以下命令。

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

然后只需在神经网络代码中添加几行即可利用自动混合精度（AMP）。在下面加粗了以下几行：

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()

在这里可以看到初始化模型，amp.也可以通过使用指定的损失比例amp.scale_loss

标杆管理

可以使用这个很棒的存储库对放大器的性能进行基准测试，该存储库对CIFAR数据集上的VGG16模型进行基准测试。只需要更改几行代码即可为我们工作。可以在此处找到修改后的版本。要自己运行基准测试代码，可能需要：

https://github.com/znxlwm/pytorch-apex-experiment

https://github.com/MLWhiz/data_science_blogs/tree/master/amp

git clone https://github.com/MLWhiz/data_science_blogs
cd data_science_blogs/amp/pytorch-apex-experiment/
python run_benchmark.py
python make_plot.py --GPU 'RTX' --method 'FP32' 'FP16' 'amp' --batch 128 256 512 1024 2048

这将在主目录中为您填充以下图形：

在这里，使用各种精度和批处理大小设置训练了同一模型的多个实例。可以看到，从FP32到安培，内存需求降低了，而精度却保持大致相同。时间也会减少，但不会减少那么多。这可能归因于简单的数据集或简单的模型。

根据NVIDIA提供的基准，自动混合精度的运行速度比标准FP32型号快3倍，如下所示。