PyTorch 性能调优指南 | （一）一般优化方法

在优化PyTorch代码性能的过程中，发现了官网的这篇 Performance Tuning Guide¹，觉得很有帮助。因此，在这一边翻译一边加入个人的实践经验供参考。

性能调优指南是一套优化方法和最佳实践方式，可以加速 PyTorch 深度学习模型的训练和推理 。所介绍的技术通常只需改变几行代码就可以实现，并且可以应用于所有领域的各种深度学习模型。

（一）一般优化方法

1.1 启用异步数据加载和扩充

torch.utils.data.DataLoader支持独立工作子进程中的异步数据加载和数据扩充。DataLoader的默认设置为num_workers=0。这意味着数据加载是同步的，在主进程中完成。因此，主训练过程必须等待数据可用后才能继续执行。

设置num_workers>0，可以支持异步数据加载以及训练与数据加载之间的重叠进行。num_workers应该根据工作负载、CPU、GPU 和训练数据的位置进行调优。

DataLoader接受pin_memory参数，其默认值为False。当使用 GPU 时，最好设置pin_memory=True，指示DataLoader使用固定内存，使内存更快地从主机异步复制到GPU上。

1.2 禁用验证或推理的梯度计算

PyTorch 会从需要梯度的张量的所有运算中保存中间缓冲区。通常，验证或推理不需要梯度。torch.no_grad()上下文管理器可以用来禁用指定代码块内的梯度计算。这将加快运行的速度并减少所需的内存。torch.no_grad()也可以用作函数的装饰器。

1.3 禁用批归一化后直连的卷积的偏置

torch.nn.Conv2d()有bias参数，默认值为True（Conv1d和Conv3d也是如此）。

如果一个nn.Conv2d层后面直接连了一个nn.BatchNorm2d层，那么就不需要卷积中的偏置，改用nn.Conv2d(..., bias=False, ....)。不需要偏置是因为在第一步中，BatchNorm减去了平均值，有效地抵消了偏置的影响。

只要BatchNorm（或其他归一化层）在与卷积偏置的相同维度上进行归一化，这个方法同样适用于一维和三维卷积。

torchvision 提供的模型已经实现了这种优化。

1.4 使用parameter.grad = None替代model.zero_grad()或optimizer.zero_grad()

不再调用

model.zero_grad()
# or
optimizer.zero_grad()

使梯度归零，改用以下方法：

for param in model.parameters():
    param.grad = None

第二个代码段不会将每个参数的内存清零，而且后续的反向传播将使用赋值而不是加法来存储梯度，这减少了内存运算的数量。

将梯度设置为None与设置为零的数值表现略有不同，详情请参阅文档²。

或者，从 PyTorch 1.7 开始，可以调用model或 optimizer.zero_grad(set_to_none=True)。

实测 PyTorch 2.0.0 set_to_none默认为True。

1.5 融合逐点运算

逐点运算（逐元素加法、乘法、数学函数 - sin()、cos()、sigmoid()等）可以被融合到一个内核中，以分摊内存访问时间和内核启动时间。

PyTorch JIT 可以自动融合内核。尽管可能有更多的融合机会尚未在编译器中实现，而且并非所有的设备类型都得到同等支持。（如果不了解内核融合，可以参见³）

逐点运算是受内存约束的，对于每个运算，PyTorch都会启动一个单独的内核。每个内核从内存中加载数据，执行计算（这一步通常不会消耗太多资源），并将结果存储回内存中。

融合的算子对于多个融合逐点运算只启动一个内核，并且只向内存加载/存储一次数据。这使得 JIT 对于激活函数、优化器、自定义 RNN 单元等非常有用。

在最简单的情况下，可以通过对函数定义应用torch.jit.script装饰器来实现融合，例如：

@torch.jit.script
def fused_gelu(x):
    return x * 0.5 * (1.0 + torch.erf(x / 1.41421))

关于更多高级用例，请参考 TorchScript 文档⁴。

1.6 为计算机视觉模型启用 channels_last 内存格式

PyTorch 1.5 为卷积网络引入了对channels_last内存格式的支持。这种格式是为了与 AMP（Automatic Mixed Precision 自动混合精度）一起使用，以进一步加速带有 Tensor Cores 的卷积神经网络。

对channels_last的支持是实验性的，但它有望用于标准的计算机视觉模型（例如ResNet-50、SSD）。要将模型转换为channels_last格式，请遵循 Channels Last 内存格式教程⁵。该教程包括关于转换现有模型的部分。

1.7 检查点中间缓冲区

缓冲区检查点是一种减轻模型训练的内存容量负担的技术。它并非存储所有层的输入来计算反向传播中的上游梯度，而是存储少数层的输入，其他层则在反向传播中重新计算。内存需求的减少可以增加批处理大小，从而提高利用率。

应谨慎选择检查点的目标。最好不要存储再计算开销较小的大层的输出。目标层的例子是激活函数（如ReLU、Sigmoid、Tanh），上/下采样和积累深度较小的矩阵-向量运算。

PyTorch 支持原生的torch.utils.checkpointAPI来自动执行检查点和再计算。

如果使用 Huggingface Transformers 的模型，可以调用函数model.gradient_checkpointing_enable()来实现。

1.8 禁用调试 API

许多 PyTorch API 是用于调试的，应在常规训练运行中禁用：

• 异常检测：torch.autograd.detect_anomaly或torch.autograd.set_detect_anomaly(True)
• 分析器相关：torch.autograd.profiler.emit_nvtx，torch.autograd.profiler.profile
• 自动微分梯度检查：torch.autograd.gradcheck或torch.autograd.gradgradcheck

---

1. https://pytorch.org/tutorials/recipes/recipes/tuning_guide.html Szymon Migacz ↩︎

2. https://pytorch.org/docs/stable/generated/torch.optim.Optimizer.zero_grad.html#torch.optim.Optimizer.zero_grad ↩︎

3. 内核融合：GPU深度学习的“加速神器” ↩︎

4. https://pytorch.org/docs/stable/jit.html ↩︎

5. https://pytorch.org/tutorials/intermediate/memory_format_tutorial.html ↩︎