备注

点击此处下载完整示例代码

性能调优指南

创建于：2025 年 4 月 1 日 | 最后更新：2025 年 4 月 1 日 | 最后验证：2024 年 11 月 5 日

作者：Szymon Migacz

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

通用优化

启用异步数据加载和增强

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

设置 num_workers > 0 启用异步数据加载和训练与数据加载之间的重叠。 num_workers 应根据工作负载、CPU、GPU 和训练数据的位置进行调整。

接受 pin_memory 个参数，默认为 False 。当使用 GPU 时，最好设置 pin_memory=True ，这指示 DataLoader 使用固定内存，并启用从主机到 GPU 的更快和异步内存复制。

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

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

直接跟随批归一化的卷积中禁用偏差

torch.nn.Conv2d() 具有默认值为 True 的 bias 参数（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 与设置为零的数值行为略有不同，更多详情请参阅文档。

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

融合操作 ¶

点操作，如逐元素加法、乘法以及数学函数如 sin()、cos()、sigmoid()等，可以合并成一个单独的内核。这种融合有助于减少内存访问和内核启动时间。通常，点操作是内存受限的；PyTorch 的 eager 模式为每个操作启动一个单独的内核，这涉及到从内存中加载数据，执行操作（通常不是最耗时的步骤），并将结果写回内存。

通过使用融合算子，对于多个逐点操作，只需启动一个内核，数据只需加载和存储一次。这种效率对于激活函数、优化器和自定义 RNN 单元等特别有益。

PyTorch 2 引入了由 TorchInductor 提供的编译模式，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 1.5 引入了对 channels_last 内存格式的支持。此格式旨在与 AMP 结合使用，以进一步加速具有 Tensor Cores 的卷积神经网络

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

检查点中间缓冲区

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

检查点目标应仔细选择。最好是不要存储具有小重新计算成本的较大层输出。示例目标层包括激活函数（例如 ReLU ， Sigmoid ， Tanh ），上/下采样以及具有小累积深度的矩阵-向量运算。

PyTorch 支持原生的 torch.utils.checkpoint API，用于自动执行检查点和重新计算。

禁用调试 API

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

异常检测：torch.autograd.detect_anomaly 或 torch.autograd.set_detect_anomaly(True)
分析器相关：torch.autograd.profiler.emit_nvtx, torch.autograd.profiler.profile
autograd gradcheck ：torch.autograd.gradcheck 或 torch.autograd.gradgradcheck

CPU 特定优化

利用非一致性内存访问（NUMA）控制

NUMA（非一致性内存访问）是一种在数据中心机器中使用的内存布局设计，旨在利用多插槽机器中具有多个内存控制器和块的内存局部性。一般来说，所有深度学习工作负载，无论是训练还是推理，在不访问 NUMA 节点之间的硬件资源的情况下，都能获得更好的性能。因此，可以通过运行多个实例来执行推理，每个实例运行在一个插槽上，以提高吞吐量。对于单节点上的训练任务，建议使用分布式训练，以便每个训练过程运行在一个插槽上。

通常情况下，以下命令仅在 Nth 节点上的核心上执行 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

使用 oneDNN Graph 与 TorchScript 进行推理

oneDNN Graph 可以显著提升推理性能。它将一些计算密集型操作（如卷积、矩阵乘法）与其相邻操作融合。在 PyTorch 2.0 中，它作为 beta 功能支持 Float32 & BFloat16 数据类型。oneDNN Graph 接收模型的图，并根据示例输入的形状识别操作融合的候选者。应使用示例输入对模型进行 JIT-trace。对于与示例输入具有相同形状的输入，在经过几轮预热迭代后，将观察到速度提升。下面的示例代码片段适用于 resnet50，但也可以扩展到使用自定义模型与 oneDNN Graph 一起使用。

# Only this extra line of code is required to use oneDNN Graph
torch.jit.enable_onednn_fusion(True)

使用 oneDNN Graph API 进行推理只需额外一行代码即可使用 Float32。如果您使用 oneDNN Graph，请避免调用 torch.jit.optimize_for_inference 。

# sample input should be of the same shape as expected inputs
sample_input = [torch.rand(32, 3, 224, 224)]
# Using resnet50 from torchvision in this example for illustrative purposes,
# but the line below can indeed be modified to use custom models as well.
model = getattr(torchvision.models, "resnet50")().eval()
# Tracing the model with example input
traced_model = torch.jit.trace(model, sample_input)
# Invoking torch.jit.freeze
traced_model = torch.jit.freeze(traced_model)

一旦使用示例输入对模型进行 JIT-trace，然后经过几轮预热运行后，就可以用于推理。

with torch.no_grad():
    # a couple of warm-up runs
    traced_model(*sample_input)
    traced_model(*sample_input)
    # speedup would be observed after warm-up runs
    traced_model(*sample_input)

虽然 oneDNN Graph 的 JIT fuser 也支持使用 BFloat16 数据类型进行推理，但只有具备 AVX512_BF16 指令集架构（ISA）的机器才能展现出 oneDNN Graph 的性能优势。以下代码片段是使用 BFloat16 数据类型进行 oneDNN Graph 推理的示例：

# AMP for JIT mode is enabled by default, and is divergent with its eager mode counterpart
torch._C._jit_set_autocast_mode(False)

with torch.no_grad(), torch.cpu.amp.autocast(cache_enabled=False, dtype=torch.bfloat16):
    # Conv-BatchNorm folding for CNN-based Vision Models should be done with ``torch.fx.experimental.optimization.fuse`` when AMP is used
    import torch.fx.experimental.optimization as optimization
    # Please note that optimization.fuse need not be called when AMP is not used
    model = optimization.fuse(model)
    model = torch.jit.trace(model, (example_input))
    model = torch.jit.freeze(model)
    # a couple of warm-up runs
    model(example_input)
    model(example_input)
    # speedup would be observed in subsequent runs.
    model(example_input)

使用 PyTorch ``DistributedDataParallel``(DDP)功能在 CPU 上训练模型

对于小型模型或内存受限模型，如 DLRM，在 CPU 上训练也是一个不错的选择。在具有多个处理器的机器上，分布式训练可以带来高效的硬件资源利用率，从而加速训练过程。经过 Intel(R) oneCCL （集体通信库）优化的 Torch-ccl，实现了 allreduce 、 allgather 、 alltoall 等集体操作，实现了 PyTorch C10D ProcessGroup API，并可以动态加载为外部 ProcessGroup 。在 PyTorch DDP 模块中实现的优化中， torch-ccl 加速了通信操作。除了对通信内核的优化外， torch-ccl 还具备计算-通信同时进行的功能。

GPU 特定优化

启用张量核心

张量核心是一种专门用于计算矩阵-矩阵乘法操作的硬件，主要用于深度学习和人工智能工作负载。张量核心具有特定的精度要求，可以通过手动或通过自动混合精度 API 进行调整。

尤其是张量操作可以利用低精度工作负载。这可以通过 torch.set_float32_matmul_precision 进行控制。默认格式设置为‘最高’，它使用张量数据类型。然而，PyTorch 提供了其他精度设置：‘高’和‘中’。这些选项优先考虑计算速度而不是数值精度。

使用 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 自动调优器：

torch.backends.cudnn.benchmark = True

自动调优器的决策可能是不确定的；不同的运行可能选择不同的算法。有关更多详细信息，请参阅 PyTorch：可重现性
在某些罕见情况下，例如输入大小高度可变时，最好禁用自动调优器来运行卷积网络，以避免与每个输入大小相关的算法选择开销。

避免不必要的 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()

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

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

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

使用混合精度和 AMP

混合精度利用 Tensor Cores，在 Volta 及更新的 GPU 架构上提供高达 3 倍的总体加速。要使用 Tensor Cores，需要启用 AMP，并且矩阵/张量维度需要满足调用使用 Tensor Cores 的内核的要求。

要使用 Tensor Cores：

设置大小为 8 的倍数（以映射到 Tensor Cores 的维度）
- 有关更多详细信息以及特定于层类型的指南，请参阅深度学习性能文档。
- 如果层的大小是从其他参数推导出来的而不是固定的，它仍然可以进行显式填充，例如 NLP 模型中的词汇表大小
启用 AMP
- 混合精度训练和 AMP 简介：视频、幻灯片
- 从 PyTorch 1.6 开始支持原生 PyTorch AMP：文档、示例、教程

预分配内存以应对可变输入长度

语音识别或 NLP 模型通常在具有可变序列长度的输入张量上进行训练。可变长度可能会对 PyTorch 缓存分配器造成问题，导致性能降低或出现意外的内存不足错误。如果短序列长度的批次后面紧接着是长序列长度的批次，那么 PyTorch 被迫释放之前迭代中的中间缓冲区，并重新分配新的缓冲区。这个过程耗时且会导致缓存分配器碎片化，可能引发内存不足错误。

一种典型的解决方案是实现预分配。它包括以下步骤：

生成一个（通常是随机生成的）具有最大序列长度的输入批次（对应于训练数据集中的最大长度或某个预定义的阈值）
执行一个前向和后向的批次，不要执行优化器或学习率调度器，此步骤预分配最大大小的缓冲区，这些缓冲区可以在后续训练迭代中重复使用
将梯度置零
继续进行常规训练

分布式优化

使用高效的数据并行后端

PyTorch 有两种方式实现数据并行训练：

提供了更好的性能和扩展到多个 GPU 的能力。更多信息请参阅 PyTorch 文档中 CUDA 最佳实践的相应部分。

跳过使用 `DistributedDataParallel` 和梯度累积时的不必要的 all-reduce 操作

默认情况下，torch.nn.parallel.DistributedDataParallel 在每个反向传播后执行梯度 all-reduce，以计算所有参与训练的工作进程的平均梯度。如果训练使用 N 步梯度累积，则不需要在每个训练步骤后进行 all-reduce，仅在最后一个反向调用之后、执行优化器之前执行 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 梯度全归约的桶。 DistributedDataParallel 将全归约与反向传播重叠。只有当给定桶中所有参数的梯度都可用时，才会异步触发特定桶的全归约。

为了最大化重叠量，模型构造函数中的顺序应大致匹配执行过程中的顺序。如果顺序不匹配，则整个桶的全归约将等待最后到达的梯度，这可能会减少反向传播和全归约之间的重叠，全归约可能会暴露出来，从而减慢训练速度。

DistributedDataParallel 与 find_unused_parameters=False （这是默认设置）依赖于在反向传播过程中遇到的操作顺序自动形成桶。使用 find_unused_parameters=False 时，无需重新排序层或参数即可达到最佳性能。

在分布式环境中负载均衡工作负载 ¶

负载不平衡通常可能发生在处理顺序数据的模型中（语音识别、翻译、语言模型等）。如果一个设备接收到的数据批次序列长度长于其他设备的序列长度，那么所有设备都将等待最后一个完成工作的工作者。反向传播在具有 DistributedDataParallel 后端的分布式环境中充当隐式同步点。

解决负载均衡问题有多种方法。核心思想是在每个全局批次中尽可能均匀地将工作负载分配给所有工作者。例如，Transformer 通过形成大约具有恒定标记数的批次（批次中序列数量可变）来解决不平衡，其他模型通过将具有相似序列长度的样本分桶或甚至按序列长度对数据集进行排序来解决不平衡。

脚本总运行时间：（0 分钟 0.000 秒）

由 Sphinx-Gallery 生成的画廊