18个必须知道的PyTorch提速秘籍：工作原理和方法

调整深度学习管道如同找到合适的齿轮组合（图片来源：Tim Mossholder）

为什么要阅读本博？

深度学习模型的训练/推理过程涉及到多个步骤。在时间和资源受限的情况下，实验迭代速度越快，越能优化模型预测的性能。本博收集整理了几个PyTorch的技巧和秘籍，以最大限度地提高内存的使用效率，最小化运行时间。为了更好地利用这些技巧，还需要了解它的工作原理。

作为开篇，本博给出了一个完整的内容清单和代码片段，方便读者跳转到需要优化的脚本，然后一个接着一个详细地展开研究。此外，还为每个秘籍技巧提供了代码片段，并对设备类型(CPU/GPU)和模型类型做出了注释。

内容清单

• 数据加载
  1. 将活跃数据移到SSD中
  2. Dataloader(dataset, num_workers=4*num_GPU)
  3. Dataloader(dataset, pin_memory=True)
• 数据操作
  4. 在运行操作的设备上将向量、矩阵、张量直接创建为 torch.Tensor

5. 避免CPU和GPU之间不必要的数据传输

6. 使torch.from_numpy(numpy_array)torch.as_tensor(others)
7. 适用于重叠数据传输时使用 tensor.to(non_blocking=True)

8. 利用PyTorchJIT将点态（元素级）操作融合到单个内核中7.适用于重叠数据传输时使用张量(非阻塞=True)

• 模型架构
  9. 将不同架构设计的尺寸设置为8的倍数(用于混合精度的16位浮点FP16)
• 训练模型
  10. 将批大小设置为8的倍数，并最大化GPU内存的使用量
  11. 前向传递使用混合精度（但不后向传递）
  12. 在优化器更新权重之前，将梯度设置为None(例如，model.zero_grad(set_to_none=True))
  13. 梯度累积：更新其他x批的权重，以模拟更大的批大小13.梯度累积：更新其他x批的权重，以模拟更大的批大小
• 推理和验证
  14. 关闭梯度计算
• 卷积神经网络（CNN）专项

15. torch.backends.cudnn.benchmark = True
16. 4D NCHW张量使用channels_last内存格式

17. 关闭批处理归一化之前的卷积层的偏差

• 分布式优化
  18.使用DistributedDataParallel 取代DataParallel

与第7、11、12、13号秘籍相关的代码片段：

高级概念

总的来说，可以通过3个关键措施来优化时间和内存的使用情况。首先，尽可能减少i/o（输入/输出），将模型管道绑定到计算（数学受限或数学绑定），而非绑定到i/o（带宽受限或内存绑定），利用GPU的专长来加速计算；第二，尽可能多地重叠进程，以节省时间；第三，最大化内存使用效率，以节省内存。通过节约内存实现更大的批大小，达到节省更多的时间的目的。优化时间之后可以加快模型开发周期，使模型性能更优。

1.将活跃数据移到SSD中

不同机器有不同的硬盘，如HHD和SSD。建议将项目中使用的活跃数据移到SSD(或具有更好i/o的硬盘驱动器)之中，以获得更快的速度。

#CPU #GPU #SaveTime

2.异步数据加载和增强

num_workers=0仅在训练之前或者训练过程完成之后才会执行数据加载。对于i/o和大数据的增强，设置num_workers>0有望加速这个过程。对于GPU来说，通过本实验发现：num_workers = 4*num_GPU 的性能最好。话虽如此，也可以测试机器最适合的num_workers。需要注意的是， num_workers越高，内存开销便越大，这完全是在意料之中，因为需要在内存中处理更多的数据副本。

#CPU #GPU #SaveTime

3. 使用固定内存来减少数据传输

设置pin_memory=True会跳过从可分页内存到固定内存的数据传输（作者提供的图片，灵感来自于此图片）

GPU不能直接从CPU的可分页内存中访问数据。设置pin_memory=True可以直接为CPU主机上的数据分配分段内存，并节省将数据从可分页存储区传输到分段内存（即固定内存，锁定分页内存）的时间。此设置可以与num_workers=4*num_GPU结合使用。

#GPU #SaveTime

4.直接将向量、矩阵、张量创建为在设备上运行的torch.Tensor

当PyTorch需要用到torch.Tensor数据的时候，首先应尝试在运行它们的设备上创建它们。不要使用本机Python或NumPy来创建数据，然后再将其转换为torch.Tensor。在大多数情况下，如果打算在GPU中使用它们，则直接在GPU中创建它们。

唯一的语法差异是，NumPy中的随机数生成需要一个额外的随机数，例如，np.random.rand()vstorch.rand()。许多其他的函数在NumPy中都有相应的与之对应的函数：

#GPU #SaveTime

5. 避免CPU和GPU之间不必要的数据传输

正如在高级概念中所述，应尽可能多地减少i/o，注意下述命令：

#GPU #SaveTime

6.使用torch.from_numpy(numpy_array)和torch.as_tensor(others)代替torch.tensor，torch.tensor()

如果源设备和目标设备都是CPU，则torch.from_numpy和torch.as_tensor不会创建数据副本。如果源数据是NumPy数组，则使用torch.from_numpy(numpy_array)会更快。如果源数据是具有相同数据类型和设备类型的张量，那么torch.as_tensor(others)可以在适用的情况下避免复制数据。others 可以是Python列表、元组或torch.tensor。如果源设备和目标设备不同，那么建议使用下一个秘籍。

#CPU #SaveTime

7.当使用重叠数据传输和内核执行时，采用tensor.to(non_blocking=True)

重叠数据传输可以减少运行时间。（由作者提供本图片）

本质上，non_blocking=True通过异步数据传输来减少执行时间。

#GPU #SaveTime

8.通过PyTorchJIT将点态（元素级）操作融合到单个内核中

包括通用数学运算在内的点态运算（参见示例列表）通常与内存绑定，PyTorchJIT自动将相邻的点态操作融合到一个内核中，以保存多次内存读写。（相当神奇，对吧？）例如，对于一个100万量级的向量，gelu函数通过将5个内核融合成1个，可以提速4倍。更多关于PyTorchJIT优化的例子可以在这里和这里找到。

9 & 10.将不同架构的批大小设置为8的倍数

为了最大限度地提高GPU的计算效率，最好确保不同的架构设计（包括神经网络的输入和输出大小/维度/通道数和批大小）是8的倍数，甚至是2的倍数（例如，64,128和256）。因为当矩阵维数对齐为2次幂的倍数时，Nvidiagpu的张量核在矩阵乘法方面获得了最优的性能。矩阵乘法是最常用的运算，也可能成为瓶颈，所以确保的张量/矩阵/向量的维数为2的幂整除数（例如，8,64,128及高达256）。

实验表明，将输出维度和批大小设置为8的倍数（即33712、4088、4096），与将输出维度和批大小设置为不能被8整除的数，（比如输出维度为33708，批大小为4084和4095）相比较，计算速度可提高1.3~4倍。提速的幅度还取决于计算类型（例如，向前通道或梯度计算）和cuBLAS版本。特别是，如果使用NLP，应检查输出的维度，通常是指词汇量大小。

使用大于256的倍数不会带来更多的好处，但也无伤大雅。输出维度和批大小设置还与cuBLAS、cuDNN版本和GPU架构相关。可以在这里找到矩阵维度贵张量核的要求。由于目前PyTorchAMP主要使用FP16， FP16为8的倍数，所以通常建议使用8的倍数。如果有一个更高级的GPU，比如A100，那么可以选择64的倍数。如果使用的是AMDGPU，则需要查阅相关AMD的文档。

除了将批大小设置为8的倍数外，还需要将批大小最大化，直到它达到GPU的内存限制。这样，就可以花更少的时间来完成一个epoch。

#GPU #SaveTime

11.在前向传递中使用混合精度，不要在反向传递中使用混合精度

有些操作不需要64位浮点或32位浮点的精度。因此，将操作设置为较低的精度可以节省内存，加快执行时间。对于各种应用，Nvidia报告使用具有张量核GPU的混合精度可以提速3.5倍到25倍。

值得注意的是，通常矩阵越大，混合精度能提速越高。在大型的神经网络(如BERT)中，实验表明，混合精度可以将训练提速2.75倍，减少37%的内存使用。具有Volta、Turing、Ampere或Hopper架构的新型GPU设备(如T4、V100、RTX 2060、2070、2080、2080Ti、A100、RTX 3090、RTX 3080和RTX 3070)可以从混合精度中获益更多，因为它们具有张量核架构，该架构具有特殊的优势，性能完胜CUDA核心。

具有张量核的NVIDIA架构支持不同的精度（图片由作者提供；数据来源）

需要注意的是，具有Hopper架构的H100，预计将在2022年第三季度发布，支持FP8（8位浮点数）。PyTorchAMP可能也会支持FP8(当前的v1.11.0还不支持FP8)。

在实际工作中，需要在模型精度性能和速度性能之间找到一个最佳平衡点。我发现混合精度的确会降低模型的性能，与算法、数据和问题也有关。

PyTorch很容易将混合精度与自动混合精度(AMP)包区别开来。PyTorch中的默认的浮点类型是32位浮点数。AMP通过使用一组16位浮点数操作节省内存和时间 (例如，matmul, linear, conv2d等，参见完整列表)。AMP自动将其转换为32位浮点数(例如，mse_loss，softmax等，参见完整列表)。某些操作（例如，add，参见完整列表）会对最宽的输入数据类型进行操作。例如，如果一个变量是32位浮点数，而另一个变量是16位浮点数，则加法结果将是32位浮点数。

自动强制转换autocast会自动将精度应用于不同的操作。由于“损失”和“梯度”是以16位浮点精度计算的，梯度可能会舍掉，当梯度值太小时便成为零。GradScaler将损失(es)乘以一个放大因子，根据放大后的损失(es)计算梯度，在优化器更新权重之前将放大后的梯度恢复回来，从而防止梯度变为零。如果 inf或NaN的缩放因子太大或太小，，那么缩放器将为下一次迭代更新缩放因子。

还可以在前向传递函数的渲染器中使用自动强制转换autocast 。

12. 在优化器更新权重之前，将梯度设置为None

通过model.zero_grad()或optimizer.zero_grad()将梯度设置为零，执行memset读写操作时会更新所有参数和梯度。但是，将梯度设置为None后不会执行memset，并且只在写入操作时更新梯度。所以，将梯度设置为None会更快一些。

13. 梯度累积：更新每个x批的权重，以模拟更大的批大小

这个技巧从更多的数据样本中积累梯度，从而使梯度的估计更为准确，使权重更加接近局部/全局最小值。当批大小很小时(由于GPU内存限制或样本的数据量很大)，这一招非常管用。

14.关闭梯度计算，以进行推理和验证

本质上，如果只需要计算模型的输出，那么在推理和验证步骤就不需要进行梯度计算。PyTorch使用一个中间内存缓冲区作为操作中间变量的要求 requires_grad=True。因此，如果已知不需要任何涉及梯度的操作，便可以在推理和验证过程中禁用梯度计算来节省资源。

15. torch.backends.cudnn.benchmark = True

在训练循环之前设置torch.backends.cudnn.benchmark=True可以加速计算。由于cuDNN算法在计算不同大小的卷积核时的性能各不相同，自动调整器通过运行一个基准测试来找到最佳的算法（目前的算法有这些、这些和这些）。当输入大小不经常改变时，建议使用打开的设置。如果输入大小经常发生变化，那么自动调整器需要过频繁地进行基准测试，这可能会影响性能。前向传播时它可以加速1.27倍，后向传播时它可以加速1.70倍。

16. 4D NCHW张量使用channels_last内存格式

4D NCHW被重新组织为NHWC格式（作者图片的灵感来自参考文献）

使用chanes_last内存格式，按像素对像素的方式保存图像，这种格式为最密集的内存格式。原始的4D NCHW张量将内存中的每个通道（红色/灰色/蓝色）聚集到一起。转换后，x=x.to(memory_format=torch.channels_last)，数据在内存中被重新组织为NHWC(channels_last 格式)。此时，RGB层的每个像素都更加接近。这种NHWC格式与AMP的16位浮点相比，可以实现8%到35%的倍速）。

目前，它仍处于测试阶段，只支持4D NCHW张量和某些模型(例如，alexnet, mnasnet family, mobilenet_v2, resnet family, shufflenet_v2, squeezenet1, vgg ，参见完整列表)。但可以肯定地看出，它将成为一项标准的优化。

17.关闭在批处理归一化之前的卷积层偏差

在数学上，偏差效应将通过批归一化的平均减法来抵消，这种方式在节省模型参数、降低运行时长和降低内存消耗三方面均非常有效。

18. 采用 DistributedDataParallel 代替 DataParallel

对于多GPU来说，即便只有一个节点，它总是更偏爱 DistributedDataParallel，因为 DistributedDataParallel采用了多进程，为每个GPU创建一个进程来绕过Python全局解释器锁 (GIL)，从而加快运行速度。

#GPU #DistributedOptimizations #SaveTime

总结

在本文中，制作了一个内容清单，并提供了18个PyTorch代码片段。然后，解释了它们的工作原理，对各个方面的工作逐一展开，内容包括数据加载、数据操作、模型架构、训练、推理、特定于cnn的优化和分布式计算。深入理解了它们的工作原理之后，便能够找到适用于任何深度学习框架中的深度学习建模的通用准则。

希望你会喜欢更加高效的PyTorch，并学习到新的知识！

如果有任何意见或建议，请在评论区留下评论或其他建议。感谢拨冗阅读。

希望我们的深度学习管道能够像火箭一样“一飞冲天”：D(图片由Bill Jelen拍摄)

衷心感谢 Katherine Prairie