注意
跳到最后下载完整示例代码。
性能调优指南#
创建于：2020 年 9 月 21 日 | 最后更新：2025 年 7 月 9 日 | 最后验证：2024 年 11 月 5 日
性能调优指南是一套优化和最佳实践，可以加速 PyTorch 中深度学习模型的训练和推理。所提出的技术通常只需更改几行代码即可实现，并且可以应用于所有领域的各种深度学习模型。
PyTorch 模型的通用优化技术
CPU 专用性能优化
GPU 加速策略
分布式训练优化
PyTorch 2.0 或更高版本
Python 3.8 或更高版本
支持 CUDA 的 GPU（建议用于 GPU 优化）
Linux、macOS 或 Windows 操作系统
概述#
性能优化对于高效的深度学习模型训练和推理至关重要。本教程涵盖了一套全面的技术，用于加速不同硬件配置和用例下的 PyTorch 工作负载。
通用优化#
import torch
import torchvision
启用异步数据加载和增强#
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 的更快异步内存复制。
禁用验证或推理的梯度计算#
PyTorch 会保存所有涉及需要梯度的张量的操作的中间缓冲区。通常，验证或推理不需要梯度。torch.no_grad() 上下文管理器可用于在指定的代码块内禁用梯度计算，这会加速执行并减少所需的内存量。torch.no_grad() 也可以用作函数装饰器。
禁用紧随批归一化层之后的卷积的偏差#
torch.nn.Conv2d() 有一个 bias 参数，默认为 True（Conv1d 和 Conv3d 也是如此）。
如果 nn.Conv2d 层紧随 nn.BatchNorm2d 层之后，则卷积中的偏差是不必要的，此时应使用 nn.Conv2d(..., bias=False, ....)。偏差是不必要的，因为在第一步中，BatchNorm 会减去均值，这有效地抵消了偏差的影响。
这也适用于 1d 和 3d 卷积，只要 BatchNorm（或其他归一化层）在与卷积偏差相同的维度上进行归一化。
从 torchvision 获得的模型已经实现了此优化。
使用 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 与将其设置为零相比，在数值行为上略有不同，更多详情请参阅文档。
或者，调用 model 或 optimizer.zero_grad(set_to_none=True)。
融合操作#
逐点操作，例如逐元素加法、乘法，以及 sin()、cos()、sigmoid() 等数学函数，可以合并到一个内核中。这种融合有助于减少内存访问和内核启动时间。通常，逐点操作受内存限制；PyTorch 急切模式为每个操作启动一个单独的内核，这涉及从内存加载数据、执行操作（通常不是最耗时的步骤）以及将结果写回内存。
通过使用融合运算符，对于多个逐点操作，只需启动一个内核，并且数据只加载和存储一次。这种效率对于激活函数、优化器和自定义 RNN 单元等尤其有益。
PyTorch 2 引入了由 TorchInductor（底层编译器）实现的编译模式，该编译器可自动融合内核。TorchInductor 的功能不仅限于简单的逐元素操作，还能实现符合条件的逐点和归约操作的高级融合，从而优化性能。
最简单的情况下，可以通过将 torch.compile 装饰器应用于函数定义来启用融合，例如
@torch.compile
def gelu(x):
return x * 0.5 * (1.0 + torch.erf(x / 1.41421))
有关更高级的用例，请参阅torch.compile 简介。
为计算机视觉模型启用 channels_last 内存格式#
PyTorch 支持卷积网络的 channels_last 内存格式。此格式旨在与 AMP 结合使用，以进一步加速使用 Tensor Cores 的卷积神经网络。
对 channels_last 的支持仍处于实验阶段，但预计它适用于标准计算机视觉模型（例如 ResNet-50、SSD）。要将模型转换为 channels_last 格式，请遵循Channels Last 内存格式教程。该教程包括一个关于转换现有模型的部分。
检查点中间缓冲区#
缓冲区检查点是一种减轻模型训练内存容量负担的技术。它不存储所有层的输入以在反向传播中计算上游梯度，而是存储少量层的输入，而其他层在反向传播期间重新计算。减少的内存需求能够增加批次大小，从而提高利用率。
检查点目标应仔细选择。最好不要存储计算成本低但输出较大的层。示例目标层包括激活函数（例如 ReLU、Sigmoid、Tanh）、上/下采样以及累积深度较小的矩阵-向量操作。
PyTorch 支持原生的 torch.utils.checkpoint API，可自动执行检查点和重新计算。
禁用调试 API#
许多 PyTorch API 旨在用于调试，应在常规训练运行中禁用
CPU 专用优化#
利用非统一内存访问 (NUMA) 控制#
NUMA，即非统一内存访问，是数据中心机器中使用的内存布局设计，旨在利用多插槽机器中内存的局部性，这些机器具有多个内存控制器和块。总的来说，所有深度学习工作负载（训练或推理）在不跨 NUMA 节点访问硬件资源的情况下都能获得更好的性能。因此，推理可以运行多个实例，每个实例在一个插槽上运行，以提高吞吐量。对于单节点上的训练任务，建议使用分布式训练，使每个训练进程在一个插槽上运行。
在一般情况下，以下命令仅在第 N 个节点上的核心上执行 PyTorch 脚本，并避免跨插槽内存访问以减少内存访问开销。
numactl --cpunodebind=N --membind=N python <pytorch_script>
更详细的描述可以在此处找到。
利用 OpenMP#
OpenMP 用于为并行计算任务带来更好的性能。OMP_NUM_THREADS 是最简单的加速计算开关。它决定了用于 OpenMP 计算的线程数。CPU 亲和性设置控制工作负载在多个核心上的分布方式。它影响通信开销、缓存行失效开销或页面抖动，因此正确的 CPU 亲和性设置会带来性能优势。GOMP_CPU_AFFINITY 或 KMP_AFFINITY 决定如何将 OpenMP* 线程绑定到物理处理单元。详细信息可以在此处找到。
使用以下命令，PyTorch 将在 N 个 OpenMP 线程上运行任务。
export OMP_NUM_THREADS=N
通常，以下环境变量用于设置 GNU OpenMP 实现的 CPU 亲和性。OMP_PROC_BIND 指定线程是否可以在处理器之间移动。将其设置为 CLOSE 会使 OpenMP 线程靠近主线程，位于连续的位置分区。OMP_SCHEDULE 决定 OpenMP 线程的调度方式。GOMP_CPU_AFFINITY 将线程绑定到特定的 CPU。一个重要的调优参数是核心绑定，它可以防止线程在多个 CPU 之间迁移，从而增强数据局部性并最大限度地减少核心间通信。
export OMP_SCHEDULE=STATIC
export OMP_PROC_BIND=CLOSE
export GOMP_CPU_AFFINITY="N-M"
Intel OpenMP 运行时库 (libiomp)#
默认情况下，PyTorch 使用 GNU OpenMP (GNU libgomp) 进行并行计算。在 Intel 平台上，Intel OpenMP 运行时库 (libiomp) 提供 OpenMP API 规范支持。它有时会比 libgomp 带来更多的性能优势。利用环境变量 LD_PRELOAD 可以将 OpenMP 库切换到 libiomp
export LD_PRELOAD=<path>/libiomp5.so:$LD_PRELOAD
与 GNU OpenMP 中的 CPU 亲和性设置类似，libiomp 中提供了环境变量来控制 CPU 亲和性设置。KMP_AFFINITY 将 OpenMP 线程绑定到物理处理单元。KMP_BLOCKTIME 设置线程在完成并行区域执行后，进入休眠状态前应等待的时间（以毫秒为单位）。在大多数情况下，将 KMP_BLOCKTIME 设置为 1 或 0 可获得良好的性能。以下命令显示了 Intel OpenMP 运行时库的常见设置。
export KMP_AFFINITY=granularity=fine,compact,1,0
export KMP_BLOCKTIME=1
切换内存分配器#
对于深度学习工作负载，Jemalloc 或 TCMalloc 通过尽可能重用内存，可以比默认的 malloc 函数获得更好的性能。Jemalloc 是一个通用 malloc 实现，强调避免碎片和可伸缩并发支持。TCMalloc 还具有一些优化功能，可以加速程序执行。其中之一是在缓存中保留内存，以加速常用对象的访问。即使在释放后仍保留此类缓存也有助于避免昂贵的系统调用，如果这些内存稍后被重新分配的话。使用环境变量 LD_PRELOAD 来利用其中一个。
export LD_PRELOAD=<jemalloc.so/tcmalloc.so>:$LD_PRELOAD
使用 PyTorch ``DistributedDataParallel``(DDP) 功能在 CPU 上训练模型#
对于小规模模型或内存密集型模型（如 DLRM），在 CPU 上训练也是一个不错的选择。在具有多个插槽的机器上，分布式训练可以高效利用硬件资源，加速训练过程。Torch-ccl 针对 Intel® oneCCL（集体通信库）进行了优化，可实现高效的分布式深度学习训练，实现诸如 allreduce、allgather、alltoall 等集体操作，并实现了 PyTorch C10D ProcessGroup API，可以作为外部 ProcessGroup 动态加载。除了 PyTorch DDP 模块中实现的优化之外，torch-ccl 还加速了通信操作。除了对通信内核进行的优化之外，torch-ccl 还具有同时进行计算-通信的功能。
GPU 专用优化#
启用 Tensor 核心#
Tensor 核心是专门用于计算矩阵-矩阵乘法操作的硬件，主要用于深度学习和 AI 工作负载。Tensor 核心有特定的精度要求，可以通过手动或通过自动混合精度 API 进行调整。
特别是，张量操作利用较低精度的工作负载。这可以通过 torch.set_float32_matmul_precision 进行控制。默认格式设置为“highest”，它使用张量数据类型。但是，PyTorch 提供了替代的精度设置：“high”和“medium”。这些选项优先考虑计算速度而非数值精度。
使用 CUDA 图#
在使用 GPU 时，工作首先必须从 CPU 启动，在某些情况下，CPU 和 GPU 之间的上下文切换可能导致资源利用率低下。CUDA 图是一种将计算保持在 GPU 内而无需支付内核启动和主机同步额外成本的方法。
# It can be enabled using
torch.compile(m, "reduce-overhead")
# or
torch.compile(m, "max-autotune")
CUDA 图的支持正在开发中，它的使用可能会增加设备内存消耗，并且某些模型可能无法编译。
启用 cuDNN 自动调优器#
NVIDIA cuDNN 支持多种算法来计算卷积。自动调优器会运行一个简短的基准测试，并根据给定的输入大小，为给定的硬件选择性能最佳的内核。
对于卷积网络（其他类型目前不支持），在启动训练循环之前，通过设置以下内容启用 cuDNN 自动调优器
自动调优器的决策可能是不确定的；不同的运行可能会选择不同的算法。更多详情请参阅PyTorch：可复现性
在某些罕见情况下，例如输入大小高度可变的情况下，最好禁用自动调优器来运行卷积网络，以避免与每个输入大小的算法选择相关的开销。
避免不必要的 CPU-GPU 同步#
避免不必要的同步，让 CPU 尽可能超前于加速器运行，以确保加速器工作队列包含大量操作。
尽可能避免需要同步的操作，例如
print(cuda_tensor)cuda_tensor.item()内存复制：
tensor.cuda()、cuda_tensor.cpu()和等效的tensor.to(device)调用cuda_tensor.nonzero()依赖于 CUDA 张量上操作结果的 Python 控制流，例如
if (cuda_tensor != 0).all()
直接在目标设备上创建张量#
与其调用 torch.rand(size).cuda() 来生成随机张量，不如直接在目标设备上生成输出：torch.rand(size, device='cuda')。
这适用于所有创建新张量并接受 device 参数的函数：torch.rand()、torch.zeros()、torch.full() 等。
使用混合精度和 AMP#
混合精度利用 Tensor 核心，并在 Volta 及更新的 GPU 架构上提供高达 3 倍的整体加速。要使用 Tensor 核心，应启用 AMP，并且矩阵/张量维度应满足调用使用 Tensor 核心的内核的要求。
要使用 Tensor 核心
在输入长度可变的情况下预分配内存#
用于语音识别或 NLP 的模型通常在具有可变序列长度的输入张量上进行训练。可变长度可能对 PyTorch 缓存分配器造成问题，并可能导致性能下降或意外的内存不足错误。如果一个短序列长度的批次之后是另一个长序列长度的批次，那么 PyTorch 会被迫释放前一次迭代的中间缓冲区并重新分配新缓冲区。这个过程耗时，并导致缓存分配器中的碎片化，这可能导致内存不足错误。
一个典型的解决方案是实现预分配。它包括以下步骤：
生成一个（通常是随机的）具有最大序列长度的输入批次（对应于训练数据集中的最大长度或一些预定义的阈值）
使用生成的批次执行前向和后向传播，不要执行优化器或学习率调度器，此步骤预分配最大大小的缓冲区，这些缓冲区可以在随后的训练迭代中重复使用
清零梯度
继续常规训练
分布式优化#
使用高效的数据并行后端#
PyTorch 有两种实现数据并行训练的方式
DistributedDataParallel 提供了更好的性能和扩展到多个 GPU。更多信息请参阅 PyTorch 文档中 CUDA 最佳实践的相关部分。
如果使用 DistributedDataParallel 和梯度累积进行训练，则跳过不必要的 all-reduce#
默认情况下，torch.nn.parallel.DistributedDataParallel 在每次反向传播后执行梯度 all-reduce，以计算参与训练的所有工作器上的平均梯度。如果训练使用 N 步的梯度累积，则不需要在每个训练步骤后都进行 all-reduce，只需在最后一次调用 backward 之后，即在执行优化器之前执行 all-reduce。
DistributedDataParallel 提供了 no_sync() 上下文管理器，它会禁用特定迭代的梯度 all-reduce。no_sync() 应该应用于梯度累积的前 N-1 次迭代，最后一次迭代应遵循默认执行并执行所需的梯度 all-reduce。
如果使用 DistributedDataParallel(find_unused_parameters=True)，则在构造函数和执行期间匹配层顺序#
torch.nn.parallel.DistributedDataParallel 与 find_unused_parameters=True 一起使用时，会使用模型构造函数中的层和参数顺序来为 DistributedDataParallel 梯度 all-reduce 构建桶。DistributedDataParallel 会将 all-reduce 与反向传播重叠。只有当给定桶中所有参数的梯度都可用时，才会异步触发特定桶的 all-reduce。
为了最大限度地提高重叠量，模型构造函数中的顺序应大致与执行期间的顺序匹配。如果顺序不匹配，则整个桶的 all-reduce 会等待最后到达的梯度，这可能会减少反向传播和 all-reduce 之间的重叠，all-reduce 最终可能会暴露出来，从而减慢训练速度。
DistributedDataParallel 与 find_unused_parameters=False（这是默认设置）依赖于根据反向传播过程中遇到的操作顺序自动形成桶。使用 find_unused_parameters=False 时，无需重新排序层或参数即可实现最佳性能。
在分布式设置中进行负载均衡#
负载不平衡通常可能发生在处理顺序数据（语音识别、翻译、语言模型等）的模型中。如果一个设备接收到的数据批次序列长度长于其余设备的序列长度，那么所有设备都会等待最后一个完成的工作器。反向传播在具有 DistributedDataParallel 后端的分布式设置中充当隐式同步点。
有多种方法可以解决负载平衡问题。核心思想是在每个全局批次中尽可能均匀地将工作负载分配给所有工作器。例如，Transformer 通过形成具有近似恒定数量的 token（以及批次中可变数量的序列）的批次来解决不平衡问题，其他模型通过将具有相似序列长度的样本分桶甚至通过按序列长度排序数据集来解决不平衡问题。
结论#
本教程涵盖了一套全面的 PyTorch 模型性能优化技术。主要要点包括
通用优化：启用异步数据加载，禁用推理梯度，使用
torch.compile融合操作，并使用高效的内存格式CPU 优化：利用 NUMA 控制，优化 OpenMP 设置，并使用高效的内存分配器
GPU 优化：启用 Tensor 核心，使用 CUDA 图，启用 cuDNN 自动调优器，并实现混合精度训练
分布式优化：使用 DistributedDataParallel，优化梯度同步，并平衡设备上的工作负载
这些优化中的许多都可以通过最少的代码更改来实现，并在各种深度学习模型中提供显著的性能改进。
