注意
点击 此处 下载完整的示例代码
性能调优指南¶
作者: Szymon Migacz
性能调优指南是一套优化和最佳实践，可以加速 PyTorch 中深度学习模型的训练和推理。所介绍的技术通常只需更改几行代码即可实现，并且可以应用于所有领域的各种深度学习模型。
通用优化¶
启用异步数据加载和增强¶
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 减去均值，这有效地抵消了偏置的影响。
只要 BatchNorm（或其他归一化层）在与卷积的偏置相同的维度上进行归一化，这也适用于 1D 和 3D 卷积。
来自 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 急切模式为每个操作启动一个单独的内核，这涉及从内存加载数据、执行操作（通常不是最耗时的步骤）并将结果写回内存。
通过使用融合运算符，仅为多个逐元素操作启动一个内核，并且数据仅加载和存储一次。这种效率对于激活函数、优化器和自定义 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 格式，请遵循 通道最后内存格式教程。本教程包含关于 转换现有模型 的部分。
检查点中间缓冲区¶
缓冲区检查点是一种缓解模型训练内存容量负担的技术。它不是存储所有层的输入以在反向传播中计算上游梯度，而是存储一些层的输入，并在反向传播期间重新计算其他层。减少的内存需求能够增加批大小，从而提高利用率。
应仔细选择检查点目标。最好不要存储具有较小重新计算成本的大型层输出。示例目标层是激活函数（例如 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
使用 oneDNN Graph 与 TorchScript 进行推理¶
oneDNN Graph 可以显著提升推理性能。它融合了一些计算密集型操作，例如卷积、矩阵乘法及其相邻操作。在 PyTorch 2.0 中，它作为 Float32 和 BFloat16 数据类型的 Beta 功能得到支持。oneDNN Graph 接收模型的图并根据示例输入的形状识别算子融合的候选者。模型应该使用示例输入进行 JIT 追踪。在针对与示例输入形状相同的输入进行几次预热迭代后，将观察到加速效果。以下代码片段示例适用于 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 追踪，则可以在几次预热运行后将其用于推理。
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 融合器也支持使用 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 上进行训练也是一个不错的选择。在具有多个插槽的机器上，分布式训练可以带来高效的硬件资源利用率，从而加速训练过程。 Torch-ccl 使用英特尔(R) oneCCL（集体通信库）进行了优化，用于高效的分布式深度学习训练，实现了诸如 allreduce、allgather、alltoall 等集体通信操作，实现了 PyTorch C10D ProcessGroup API，并且可以作为外部 ProcessGroup 动态加载。在 PyTorch DDP 模块中实现的优化基础上，torch-ccl 加速了通信操作。除了对通信内核进行优化外，torch-ccl 还具有同时计算通信的功能。
GPU 特定优化¶
启用 Tensor Core¶
Tensor Core 是专门为计算矩阵乘法运算而设计的硬件，主要用于深度学习和 AI 工作负载。Tensor Core 具有特定的精度要求，可以通过手动或通过自动混合精度 API 进行调整。
特别是，张量操作利用了较低精度的计算负载。可以通过 torch.set_float32_matmul_precision 进行控制。默认格式设置为“最高”，它使用张量数据类型。但是，PyTorch 提供了其他精度设置：“高”和“中”。这些选项优先考虑计算速度而不是数值精度。”
使用 CUDA Graph¶
在使用 GPU 时，首先必须从 CPU 启动工作，在某些情况下，CPU 和 GPU 之间的上下文切换会导致资源利用率下降。CUDA Graph 是一种在 GPU 内部保持计算而无需支付内核启动和主机同步额外成本的方法。
# It can be enabled using
torch.compile(m, "reduce-overhead")
# or
torch.compile(m, "max-autotune")
对 CUDA Graph 的支持正在开发中，其使用可能会导致设备内存消耗增加，并且某些模型可能无法编译。
启用 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()依赖于在 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 Core，并在 Volta 及更新的 GPU 架构上提供了高达 3 倍的整体加速。要使用 Tensor Core，应启用 AMP，并且矩阵/张量的维度应满足调用使用 Tensor Core 的内核的要求。
要使用 Tensor Core
在输入长度可变的情况下预分配内存¶
语音识别或 NLP 模型通常在输入张量上进行训练，这些张量的序列长度可变。可变长度对于 PyTorch 缓存分配器来说可能存在问题，并可能导致性能下降或意外的内存不足错误。如果短序列长度的批次后面跟着另一个长序列长度的批次，则 PyTorch 将被迫释放前一次迭代的中间缓冲区并重新分配新的缓冲区。此过程非常耗时，并导致缓存分配器中的碎片化，这可能导致内存不足错误。
一个典型的解决方案是实现预分配。它包括以下步骤
生成具有最大序列长度的输入批次（通常是随机的）（对应于训练数据集中的最大长度或某个预定义阈值）
使用生成的批次执行前向和后向传递，不要执行优化器或学习率调度器，此步骤预分配最大大小的缓冲区，这些缓冲区可以在后续训练迭代中重复使用
将梯度清零
继续进行常规训练
分布式优化¶
使用高效的数据并行后端¶
PyTorch 有两种方法可以实现数据并行训练
DistributedDataParallel 提供了更好的性能和扩展到多个 GPU 的能力。有关更多信息，请参阅 PyTorch 文档中 CUDA 最佳实践的相关部分。
如果使用 DistributedDataParallel 和梯度累积进行训练，则跳过不必要的全减少¶
默认情况下， torch.nn.parallel.DistributedDataParallel 在每次反向传递后执行梯度全减少，以计算参与训练的所有工作程序上的平均梯度。如果训练使用 N 步梯度累积，则每次训练步骤后都不需要全减少，只需要在最后一次调用反向传播后（就在执行优化器之前）执行全减少。
DistributedDataParallel 提供了 no_sync() 上下文管理器，用于禁用特定迭代的梯度全减少。 no_sync() 应应用于梯度累积的前 N-1 次迭代，最后一次迭代应遵循默认执行并执行所需的梯度全减少。
如果使用 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 秒）
