更快的计算,更高的内存效率: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

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



​ 在这里,我们使用各种精度和批处理大小训练了同一模型。 我们可以看到,从FP32到amp,内存需求降低了,而精度却保持大致相同。时间也会减少,但不会减少那么多。 这可能因为数据集或模型比较简单。

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

来源:加速比是指单精度和自动混合精度训练固定次数下的时间的比例。

deephub翻译组:孟翔杰

已标记关键词 清除标记
©️2020 CSDN 皮肤主题: 技术工厂 设计师:CSDN官方博客 返回首页