AMP Tensor Cores节省内存PyTorch模型详解

导读

只需要添加几行代码，就可以得到更快速，更省显存的PyTorch模型。

你知道吗，在1986年Geoffrey Hinton就在Nature论文中给出了反向传播算法？

此外，卷积网络最早是由Yann le cun在1998年提出的，用于数字分类，他使用了一个卷积层。但是直到2012年晚些时候，Alexnet才通过使用多个卷积层来实现最先进的imagenet。

那么，是什么让他们现在如此出名，而不是之前呢？

只有在我们拥有大量计算资源的情况下，我们才能够在最近的过去试验和充分利用深度学习的潜力。

但是，我们是否已经足够好地使用了我们的计算资源呢？我们能做得更好吗？

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

什么是Tensor Cores?

根据NVIDIA的网站：

NVIDIA Turing和Volta GPUs都是由Tensor Cores驱动的，这是一项突破性的技术，提供了突破性的AI性能。Tensor Cores可以加速AI核心的大矩阵运算，在一次运算中就可以完成混合精度的矩阵乘法和累加运算。在一个NVIDIA GPU上有数百个Tensor Cores并行运行，这大大提高了吞吐量和效率。

简单地说，它们是专门的cores，非常适合特定类型的矩阵操作。

我们可以将两个FP16矩阵相乘，并将其添加到一个FP16/FP32矩阵中，从而得到一个FP16/FP32矩阵。Tensor cores支持混合精度数学，即以半精度(FP16)进行输入，以全精度(FP32)进行输出。上述类型的操作对许多深度学习任务具有内在价值，而Tensor cores为这种操作提供了专门的硬件。

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

• FP16需要更少的内存，因此更容易训练和部署大型神经网络。它还只需要较少的数据移动。

• 数学运算在降低精度的Tensor cores运行得更快。NVIDIA给出的Volta GPU的确切数字是：FP16的125 TFlops vs FP32的15.7 TFlops(8倍加速)。

但也有缺点。当我们从FP32转到FP16时，我们需要降低精度。

FP32 vs FP16: FP32 有8个指数位和23个分数位，而FP16有5个指数位和10个分数位。

但是FP32真的有必要吗？

实际上，FP16可以很好地表示大多数权重和梯度。所以存储和使用FP32是很浪费的。

那么，我们如何使用Tensor Cores？

我检查了一下我的Titan RTX GPU有576个tensor cores和4608个NVIDIA CUDA核心。但是我如何使用这些tensor cores呢？

坦白地说，NVIDIA用几行代码就能提供自动混合精度，因此使用tensor cores很简单。我们需要在代码中做两件事：

• 需要用到FP32的运算比如Softmax之类的就分配用FP32，而Conv之类的操作可以用FP16的则被自动分配用FP16。

• 使用损失缩放 为了保留小的梯度值。梯度值可能落在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.initialize初始化了我们的模型。我们还使用amp.scale_loss来指定损失缩放。

基准测试

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

这会在home目录中生成下面的图:

在这里，我使用不同的精度和批大小设置训练了同一个模型的多个实例。我们可以看到，从FP32到amp，内存需求减少，而精度保持大致相同。时间也会减少，但不会减少那么多。这可能是由于数据集或模型太简单。

根据NVIDIA给出的基准测试，AMP比标准的FP32快3倍左右，如下图所示。

在单精度和自动混合精度两种精度下，加速比为固定周期训练的时间比。

来源：https://juejin.cn/post/7157663977437134884