更快的计算，更高的内存效率：PyTorch混合精度模型AMP介绍

作者：Rahul Agarwal

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

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

那么，是什么让他们直到现在才这么出名？

只有在我们拥有大量计算资源的情况下，我们才能进行足够的实验并充分利用深度学习的潜力。

但是我们真的充分利用了现有的计算资源吗？我们可以做得更好吗？

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

什么是Tensor Cores？

根据NVIDIA网站的介绍：

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

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

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

与FP32相比FP16主要有一下两个好处。

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 Cores呢？

我检查了我的Titan RTX GPU。它拥有576个Tensor Cores以及4,608个NVIDIA CUDA内核。但是如何使用这些Tensor Cores？

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

1.将FP32所需的操作（如Softmax）分配给FP32，而将FP16可以完成的操作（如卷积）自动分配给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模型进行测试。我只需要更改几行代码即可为我们工作。您可以在此处找到修改后的版本。要自己运行测试代码，您可能需要：

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 12825651210242048

这将在主目录中为您生成以下图形：