PyTorch 性能调优指南 | （三）针对GPU的优化方法

Galaxy_Husky

已于 2023-07-06 17:00:59 修改

阅读量1.1k

点赞数 23

分类专栏： pytorch 文章标签： pytorch 人工智能 python

于 2023-06-28 18:42:16 首次发布

本文链接：https://blog.csdn.net/qq_32044697/article/details/131438936

版权

pytorch 专栏收录该内容

4 篇文章 0 订阅

订阅专栏

PyTorch Performance Tuning Guide

（三）针对GPU的优化方法

原本是应该先介绍第二部分针对 CPU 的优化方法，由于其中的方法自己都没有实践过，感觉绝大部分人也用不上，所以暂时先搁置，之后有时间的话再补上。

（三）针对GPU的优化方法

3.1 启用 cuDNN auto-tuner

NVIDIA cuDNN 支持许多计算卷积的算法。Autotuner 运行一个简短的基准测试，在给定的输入大小下，选择在给定硬件上具有最佳性能的内核。

对于卷积网络（目前不支持其他类型），在启动训练循环之前通过以下设置来启用 cuDNN auto-tuner：

torch.backends.cudnn.benchmark = True

auto-tuner 的选择可能是非确定性的；不同的运行可能会选择不同的算法。更多细节参见 ¹ （包括目前 PyTorch 关于复现性的所有方法）。反之，当需要减少训练的随机性来复现实验时，需要设置torch.backends.cudnn.benchmark = False。
在一些罕见的情况下，比如输入大小高度可变时，最好在禁用 auto-tuner 的情况下运行卷积网络，以避免为每个输入大小选择算法相关的开销。

3.2 避免不必要的 CPU-GPU 同步

避免不必要的同步，尽可能让 CPU 在加速器之前运行，以确保加速器的工作队列包含许多运算。

尽可能避免需要同步的运算，例如：

print(cuda_tensor)
cuda_tensor.item()
内存复制：tensor.cuda()，cuda_tensor.cpu()和等效tensor.to(device)的调用
cuda_tensor.nonzero()
Python 控制流，它取决于在 CUDA 张量上执行运算的结果，比如，if (cuda_tensor != 0).all()

3.3 直接在目标设备上创建张量

替代调用torch.rand(size).cuda()来生成随机张量，而是直接在目标设备上产生输出：torch.rand(size, device='cuda')。

这适用于所有创建新张量并接受device参数的函数：torch.rand()，torch.zeros()，torch.full()和类似的函数。

3.4 使用混合精度和 AMP

混合精度利用 Tensor Cores，在 Volta 和较新的 GPU 架构上提供高达3倍的整体速度提升。要使用 Tensor Cores，应启用AMP。矩阵/张量维度应满足调用那些使用 Tensor Cores 的内核的要求。

要使用 Tensor Cores：

将大小设置为8的倍数（以映射到 Tensor Cores 的维度）。
- 更多细节和针对层类型的指南，请参见²。
- 如果层的大小与其他参数有关而不是固定的，它仍然可以被显式填充，例如 NLP 模型中的词表大小。
启用 AMP
- 混合精度训练和AMP的介绍：视频、幻灯片（链接参见原文）。
- 从 PyTorch 1.6 开始提供原生的 PyTorch AMP ：文档、示例、教程（链接参见原文）。

个人使用下来的体验是 AMP 确实能节省不少显存，但是除了可接受的精度损失之外，有个非常棘手的问题是 loss 时常会出现 NaN 值，使得训练无法进行下去。具体的产生原因和解决办法可以参考³。简单来说就是在一些情况下，某些数值过大，超出了 fp16 的表示范围。问题因人而异，需要逐行 debug 去排查，可能会耗费大量的时间精力。

除了 AMP，还有一些很有效的节省显存的方法，比如梯度累积（gradient accumulation）、第一部分提到过的梯度检查点（gradient checkpoint）和 AdaFactor 优化器（解析参考⁴）。

3.5 在输入长度可变的情况下预分配内存

语音识别或 NLP 的模型经常是在序列长度可变的输入张量上训练的。可变长度可能会给 PyTorch 缓存分配器带来问题，会导致性能下降或出现意外的内存不足错误。如果一个序列长度较短的批次之后是另一个序列长度较长的批次，那么 PyTorch 就不得不释放上一次迭代的中间缓冲区并重新分配新的缓冲区。这个过程很耗时，而且会导致缓存分配器中出现碎片，从而可能导致内存不足的错误。

一个经典的解决方案是实施预分配。它包括以下步骤：