独家 | 兼顾速度和存储效率的PyTorch性能优化（2022）

调整深度学习管道如同找到合适的齿轮组合（图片来源：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将点态（元素级）操作融合到单个内核中

• 模型架构

9. 将不同架构设计的尺寸设置为8的倍数，使其适用于混合精度的16位浮点（FP16）。

• 训练模型

10. 将批大小设置为8的倍数，并最大化GPU内存的使用量

11. 前向传递使用混合精度（但不后向传递）

12. 在优化器更新权重之前，将梯度设置为None(例如，model.zero_grad(set_to_none=True))

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的硬盘驱动器)之中，以获得更快的速度。

2.异步进行数据加载和增强

设定num_workers=0，程序会仅在训练之前或者训练过程完成之后才会执行数据加载。对于i/o和大数据的增强，设置num_workers>0会有助于加速这个过程。通过本实验发现，对于GPU来说，当设定num_workers = 4*num_GPU 时，性能最好。虽然实验结果是这样，但是你也可以测试自己机器最适合的num_workers。需要注意的是， num_workers越高，内存开销便越大，这完全是在意料之中，因为更多的数据副本被存到了内存中。

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

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

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

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

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

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

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

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

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。如果源设备和目标设备不同，那么建议使用第七点。

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

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

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

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

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

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

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

实验表明，将输出维度和批大小设置为8的倍数（即33712、4088、4096）的计算速度，相对于将输出维度和批大小设置为不能被8整除的数（比如输出维度为33708，批大小为4084和4095）的计算速度而言，可提高1.3~4倍。除此之外，提速的幅度还取决于计算类型（例如，向前通道或梯度计算）和cuBLAS版本。特别是，如果你在自然语言处理领域工作，应当检查输出的维度（通常是指词汇量大小）。

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

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

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，参见完整列表），AMP会对最宽的输入数据类型进行处理。例如，如果一个变量是32位浮点数，而另一个变量是16位浮点数，则加法结果将是32位浮点数。

autocast会自动将各种精度应用于不同的操作。由于“损失”和“梯度”是以16位浮点精度计算的，梯度计算时可能会舍掉他们。这会使得梯度值太小时直接成为零。GradScaler会先将损失乘以一个放大因子，使用放大后的损失计算梯度，然后在优化器更新权重之前将放大后的梯度缩小回来，以此防止梯度变为零。如果 因缩放因子太大或太小，导致结果出现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)，从而加快运行速度。

总结

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

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

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

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

衷心感谢 Katherine Prairie

原文标题：

Optimize PyTorch Performance for Speed and Memory Efficiency (2022)

原文链接：

https://medium.com/towards-data-science/optimize-pytorch-performance-for-speed-and-memory-efficiency-2022-84f453916ea6

编辑：于腾凯

校对：林亦霖

译者简介

陈之炎，北京交通大学通信与控制工程专业毕业，获得工学硕士学位，历任长城计算机软件与系统公司工程师，大唐微电子公司工程师，现任北京吾译超群科技有限公司技术支持。目前从事智能化翻译教学系统的运营和维护，在人工智能深度学习和自然语言处理（NLP）方面积累有一定的经验。业余时间喜爱翻译创作，翻译作品主要有：IEC-ISO 7816、伊拉克石油工程项目、新财税主义宣言等等，其中中译英作品“新财税主义宣言”在GLOBAL TIMES正式发表。能够利用业余时间加入到THU 数据派平台的翻译志愿者小组，希望能和大家一起交流分享，共同进步

翻译组招募信息

工作内容：需要一颗细致的心，将选取好的外文文章翻译成流畅的中文。如果你是数据科学/统计学/计算机类的留学生，或在海外从事相关工作，或对自己外语水平有信心的朋友欢迎加入翻译小组。

你能得到：定期的翻译培训提高志愿者的翻译水平，提高对于数据科学前沿的认知，海外的朋友可以和国内技术应用发展保持联系，THU数据派产学研的背景为志愿者带来好的发展机遇。

其他福利：来自于名企的数据科学工作者，北大清华以及海外等名校学生他们都将成为你在翻译小组的伙伴。

点击文末“阅读原文”加入数据派团队~

转载须知

如需转载，请在开篇显著位置注明作者和出处（转自：数据派ID：DatapiTHU），并在文章结尾放置数据派醒目二维码。有原创标识文章，请发送【文章名称-待授权公众号名称及ID】至联系邮箱，申请白名单授权并按要求编辑。

发布后请将链接反馈至联系邮箱（见下方）。未经许可的转载以及改编者，我们将依法追究其法律责任。

点击“阅读原文”拥抱组织