注意
前往文末 下載完整的示例程式碼。
效能調優指南#
建立日期：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 支援在單獨的 worker 子程序中進行非同步資料載入和資料增強。`DataLoader` 的預設設定為 `num_workers=0`，這意味著資料載入是同步的，並在主程序中完成。因此，主訓練程序必須等待資料可用才能繼續執行。
將 `num_workers` 設定為 `> 0` 可以啟用非同步資料載入，並使訓練和資料載入之間重疊。`num_workers` 的值應根據工作負載、CPU、GPU 和訓練資料的位置進行調整。
`DataLoader` 接受 `pin_memory` 引數，預設為 `False`。在使用 GPU 時，最好將 `pin_memory` 設定為 `True`，這將指示 `DataLoader` 使用 pinned 記憶體，並實現從主機到 GPU 更快、非同步的記憶體複製。
停用驗證或推理的梯度計算#
PyTorch 會儲存所有涉及需要梯度的張量的操作的中間緩衝區。通常，驗證或推理不需要梯度。`torch.no_grad()` 上下文管理器可以應用於停用指定程式碼塊內的梯度計算，這可以加快執行速度並減少所需記憶體量。`torch.no_grad()` 也可以用作函式裝飾器。
當卷積層直接後跟 batch norm 時，停用卷積的偏置#
`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` 的數值行為與將其設定為零略有不同，更多細節請參閱[文件](https://pytorch.com.tw/docs/stable/optim.html#torch.optim.Optimizer.zero_grad)。
或者，呼叫 `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 簡介](https://pytorch.com.tw/tutorials/intermediate/torch_compile_tutorial.html)。
為計算機視覺模型啟用 `channels_last` 記憶體格式#
PyTorch 支援卷積網路的 `channels_last` 記憶體格式。該格式旨在與 AMP 結合使用，以使用 Tensor Cores 進一步加速卷積神經網路。
對 `channels_last` 的支援是實驗性的，但預計適用於標準的計算機視覺模型（例如 ResNet-50、SSD）。要將模型轉換為 `channels_last` 格式，請遵循[Channels Last Memory Format Tutorial](https://pytorch.com.tw/tutorials/intermediate/memory_format_tutorial.html)。該教程包含一個關於[轉換現有模型](https://pytorch.com.tw/tutorials/intermediate/memory_format_tutorial.html#converting-existing-models)的部分。
檢查點中間緩衝區#
緩衝區檢查點是一種減輕模型訓練記憶體容量負擔的技術。它不是在反向傳播中儲存所有層的輸入以計算上游梯度，而是儲存少量層的輸入，其餘的在反向傳播過程中重新計算。減少的記憶體需求允許增加批次大小，從而提高利用率。
應仔細選擇檢查點目標。最好的方法是不儲存計算成本很小的、但輸出很大的層。示例目標層包括啟用函式（如 `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 節點的硬體資源的情況下都能獲得更好的效能。因此，推理可以執行多個例項，每個例項執行在一個處理器上，以提高吞吐量。對於單節點上的訓練任務，建議使用分散式訓練，使每個訓練程序執行在一個處理器上。
在一般情況下，以下命令僅在第 N 個節點的核心上執行 PyTorch 指令碼，並避免跨處理器記憶體訪問以減少記憶體訪問開銷。
numactl --cpunodebind=N --membind=N python <pytorch_script>
更詳細的描述可以在[這裡](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/performance_tuning/tuning_guide.html)找到。
利用 OpenMP#
OpenMP 用於提高平行計算任務的效能。`OMP_NUM_THREADS` 是可用於加速計算的最簡單的開關。它決定了用於 OpenMP 計算的執行緒數。CPU 親和性設定控制工作負載在多個核心上的分佈。它會影響通訊開銷、快取行失效開銷或頁面抖動，因此正確設定 CPU 親和性可以帶來效能優勢。`GOMP_CPU_AFFINITY` 或 `KMP_AFFINITY` 決定如何將 OpenMP* 執行緒繫結到物理處理單元。詳細資訊可以在[這裡](https://intel.github.io/intel-extension-for-pytorch/cpu/latest/tutorials/performance_tuning/tuning_guide.html)找到。
使用以下命令，PyTorch 在 N 個 OpenMP 執行緒上執行任務。
export OMP_NUM_THREADS=N
通常，以下環境變數用於使用 GNU OpenMP 實現設定 CPU 親和性。`OMP_PROC_BIND` 指定執行緒是否可以跨處理器移動。將其設定為 `CLOSE` 會將 OpenMP 執行緒保持在連續分割槽內的主執行緒附近。`OMP_SCHEDULE` 決定 OpenMP 執行緒的排程方式。`GOMP_CPU_AFFINITY` 將執行緒繫結到特定的 CPU。一個重要的調優引數是核心固定（core pinning），它防止執行緒在多個 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](https://github.com/jemalloc/jemalloc) 是一個通用 `malloc` 實現，它側重於避免碎片和可擴充套件的併發支援。[TCMalloc](https://google.github.io/tcmalloc/overview.html) 也具有一些最佳化功能，可以加快程式執行速度。其中一項是持有記憶體快取以加快對常用物件的訪問。即使在釋放記憶體後也持有這些快取有助於避免昂貴的系統呼叫，如果這些記憶體稍後被重新分配的話。使用環境變數 `LD_PRELOAD` 來利用其中一個。
export LD_PRELOAD=<jemalloc.so/tcmalloc.so>:$LD_PRELOAD
使用 PyTorch ``DistributedDataParallel``(DDP) 功能在 CPU 上訓練模型#
對於小型模型或記憶體繫結的模型（如 DLRM），在 CPU 上訓練也是一個不錯的選擇。在具有多個處理器的機器上，分散式訓練可以實現高效的硬體資源利用，以加速訓練過程。`Torch-ccl` 針對 Intel(R) `oneCCL`（集合通訊庫）進行了最佳化，用於高效的分散式深度學習訓練，實現了 `allreduce`、`allgather`、`alltoall` 等集合通訊操作，實現了 PyTorch C10D `ProcessGroup` API，並且可以動態載入為外部 `ProcessGroup`。基於 PyTorch DDP 模組實現的最佳化，`torch-ccl` 加速了通訊操作。除了對通訊核心進行的最佳化外，`torch-ccl` 還具有同步計算-通訊功能。
GPU 特定最佳化#
啟用 Tensor Cores#
Tensor Cores 是專門用於計算矩陣-矩陣乘法運算的硬體，主要用於深度學習和人工智慧工作負載。Tensor Cores 具有特定的精度要求，可以透過手動調整或透過自動混合精度 API 進行調整。
特別是，張量運算利用了較低精度的工作負載。這可以透過 `torch.set_float32_matmul_precision` 進行控制。預設格式設定為“highest”，它利用了張量資料型別。但是，PyTorch 提供了替代的精度設定：“high”和“medium”。這些選項優先考慮計算速度而非數值精度。
使用 CUDA Graphs#
在使用 GPU 時，工作首先必須從 CPU 啟動，在某些情況下，CPU 和 GPU 之間的上下文切換可能導致資源利用率低下。CUDA Graphs 是一種在不產生額外的核心啟動和主機同步成本的情況下，在 GPU 內部進行計算的方法。
# It can be enabled using
torch.compile(m, "reduce-overhead")
# or
torch.compile(m, "max-autotune")
CUDA Graph 的支援尚在開發中，其使用可能會導致裝置記憶體消耗增加，並且某些模型可能無法編譯。
啟用 cuDNN 自動調優器#
[NVIDIA cuDNN](https://developer.nvidia.com/cudnn) 支援多種演算法來計算卷積。自動調優器執行一個簡短的基準測試，併為給定輸入大小的給定硬體選擇效能最佳的核心。
對於卷積網路（其他型別目前不支援），在啟動訓練迴圈之前，透過設定
自動調優器的決策可能是非確定性的；不同的執行可能會選擇不同的演算法。有關更多詳細資訊，請參閱[PyTorch：可復現性](https://pytorch.com.tw/docs/stable/notes/randomness.html?highlight=determinism)。
在某些罕見情況下，例如輸入大小變化很大的情況，最好停用自動調優器來運行卷積網路，以避免為每個輸入大小選擇演算法所帶來的開銷。
避免不必要的 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 Cores](https://www.nvidia.com/en-us/data-center/tensor-cores/)，並在 Volta 及更新的 GPU 架構上提供高達 3 倍的整體加速。要使用 Tensor Cores，應啟用 AMP，並且矩陣/張量維度應滿足呼叫使用 Tensor Cores 的核心的要求。
要使用 Tensor Cores：
將尺寸設定為 8 的倍數（以對映到 Tensor Cores 的維度）
有關更多詳細資訊和特定於層型別的指南，請參閱[深度學習效能文件](https://docs.nvidia.com/deeplearning/performance/index.html#optimizing-performance)。
如果層大小派生自其他引數而不是固定的，也可以顯式填充（例如 NLP 模型中的詞彙大小）。
啟用 AMP
混合精度訓練和 AMP 簡介：[幻燈片](https://nvlabs.github.io/eccv2020-mixed-precision-tutorial/files/dusan_stosic-training-neural-networks-with-tensor-cores.pdf)
提供原生 PyTorch AMP：[文件](https://pytorch.com.tw/docs/stable/amp.html)，[示例](https://pytorch.com.tw/docs/stable/notes/amp_examples.html#amp-examples)，[教程](https://pytorch.com.tw/tutorials/recipes/recipes/amp_recipe.html)
在輸入長度可變的情況下預分配記憶體#
語音識別或 NLP 模型通常在具有可變序列長度的輸入張量上進行訓練。可變長度可能會對 PyTorch 快取分配器造成問題，並可能導致效能下降或意外的記憶體不足錯誤。如果一個批次包含短序列長度，後面跟著另一個批次包含長序列長度，那麼 PyTorch 將被迫釋放前一次迭代的中間緩衝區並重新分配新的緩衝區。這個過程耗時，並導致快取分配器中的碎片，這可能導致記憶體不足錯誤。
一種典型的解決方案是實現預分配。它包括以下步驟：
生成一個（通常是隨機的）最大序列長度的輸入批次（對應於訓練資料集中的最大長度或某個預定義閾值）。
使用生成的批次執行前向和後向傳播，不要執行最佳化器或學習率排程器，此步驟將預分配最大尺寸的緩衝區，這些緩衝區可以在後續訓練迭代中重用。
清零梯度。
繼續進行常規訓練。
分散式最佳化#
使用高效的資料並行後端#
PyTorch 有兩種實現資料並行訓練的方法：
`DistributedDataParallel` 提供了更好的效能和多 GPU 擴充套件性。更多資訊請參閱 PyTorch 文件中[CUDA 最佳實踐的相關部分](https://pytorch.com.tw/docs/stable/notes/cuda.html#use-nn-parallel-distributeddataparallel-instead-of-multiprocessing-or-nn-dataparallel)。
在使用 `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 Cores，使用 CUDA Graphs，啟用 cuDNN 自動調優器，以及實現混合精度訓練。
分散式最佳化：使用 DistributedDataParallel，最佳化梯度同步，以及平衡跨裝置的負載。
其中許多最佳化可以透過最小的程式碼更改來實現，並在各種深度學習模型中提供顯著的效能提升。
