注意
點擊此處下載完整的範例程式碼
效能調整指南¶
建立於：2020 年 9 月 21 日 | 最後更新：2024 年 6 月 12 日 | 最後驗證：2024 年 11 月 05 日
效能調整指南是一組最佳化和最佳實踐，可以加速 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 會減去平均值，從而有效地消除偏置的影響。
這也適用於 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
第二段程式碼片段並未將每個個別參數的記憶體歸零，此外，後續的反向傳播過程使用賦值（assignment）而非加法（addition）來儲存梯度，這減少了記憶體操作的次數。
將梯度設定為 None 與將其設定為零相比，數值行為略有不同，詳情請參閱文件。
或者，從 PyTorch 1.7 開始，可以呼叫 model 或 optimizer.zero_grad(set_to_none=True)。
融合操作¶
逐點運算（Pointwise operations），例如元素級加法、乘法，以及數學函數如 sin()、cos()、sigmoid() 等，可以組合到單個核心（kernel）中。這種融合有助於減少記憶體存取和核心啟動時間。通常，逐點運算是記憶體受限的；PyTorch eager 模式會為每個操作啟動一個獨立的核心，這涉及從記憶體載入資料、執行操作（通常不是最耗時的步驟），以及將結果寫回記憶體。
透過使用融合運算子，僅為多個逐點運算啟動一個核心，並且資料僅載入和儲存一次。這種效率對於激活函數、優化器和自定義 RNN 單元等特別有益。
PyTorch 2 引入了一種由 TorchInductor 促進的編譯模式，TorchInductor 是一種底層編譯器，可自動融合核心。 TorchInductor 將其功能擴展到簡單的元素級操作之外，能夠對符合條件的逐點和歸約（reduction）操作進行高級融合，以實現最佳化效能。
在最簡單的情況下，可以通過將 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 Tutorial。 該教程包含一個關於轉換現有模型的部分。
檢查點中間緩衝區¶
緩衝區檢查點是一種減輕模型訓練的記憶體容量負擔的技術。 它不是儲存所有層的輸入以在反向傳播中計算上游梯度，而是儲存少數層的輸入，而其他層在反向傳播期間重新計算。 減少的記憶體需求能夠增加可以提高利用率的批次大小。
應仔細選擇檢查點目標。 最好不要儲存具有小重新計算成本的大型層輸出。 示例目標層是激活函數（例如 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。一個重要的調整參數是核心釘選 (core pinning)，它可以防止執行緒在多個 CPU 之間遷移，從而增強資料位置並最大限度地減少核心間的通訊。
export OMP_SCHEDULE=STATIC
export OMP_PROC_BIND=CLOSE
export GOMP_CPU_AFFINITY="N-M"
Intel OpenMP Runtime Library (libiomp)¶
預設情況下，PyTorch 使用 GNU OpenMP (GNU libgomp) 進行平行運算。在 Intel 平台上，Intel OpenMP Runtime Library (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 Runtime Library 的常見設定。
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 可以顯著提高推論效能。它將一些計算密集型運算（例如卷積、matmul）與它們的相鄰運算融合。在 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 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 上進行訓練也是一個不錯的選擇。在具有多個插槽的機器上，分散式訓練帶來了高效的硬體資源利用率，以加速訓練過程。Torch-ccl，使用 Intel(R) oneCCL (collective communications library) 進行最佳化，以實現高效的分散式深度學習訓練，實作諸如 allreduce、allgather、alltoall 之類的集體通訊，並將其作為外部 ProcessGroup 動態載入。 根據 PyTorch DDP 模組中實作的最佳化，torch-ccl 加速了通訊運算。 除了對通訊核心進行最佳化之外，torch-ccl 還具有同步計算-通訊功能。
GPU 特定最佳化¶
啟用 Tensor 核心¶
Tensor 核心是專門設計用於計算矩陣-矩陣乘法運算的硬體，主要用於深度學習和 AI 工作負載。 Tensor 核心具有特定的精度要求，可以手動調整或透過自動混合精度 API 進行調整。
特別是，Tensor 運算會利用較低精度的 workload。這可以透過 torch.set_float32_matmul_precision 進行控制。預設格式設定為“highest”，它會使用 Tensor 資料類型。但是，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 auto-tuner¶
NVIDIA cuDNN 支援多種演算法來計算卷積。Autotuner 會執行一個簡短的基準測試，並針對給定的輸入大小，在給定的硬體上選擇效能最佳的核心。
對於卷積網路（目前不支援其他類型），在啟動訓練迴圈之前，透過設定以下變數來啟用 cuDNN autotuner：
torch.backends.cudnn.benchmark = True
auto-tuner 的決策可能是不具決定性的；不同的執行可能會選擇不同的演算法。 更多詳細資訊請參閱 PyTorch：再現性
在某些罕見的情況下，例如輸入大小高度可變時，最好禁用 autotuner 來執行卷積網路，以避免每次輸入大小選擇演算法所帶來的額外負擔。
避免不必要的 CPU-GPU 同步¶
避免不必要的同步，讓 CPU 盡可能地在加速器之前執行，以確保加速器工作佇列包含許多操作。
在可能的情況下，避免需要同步的操作，例如：
print(cuda_tensor)cuda_tensor.item()記憶體複製：
tensor.cuda()、cuda_tensor.cpu()和等效的tensor.to(device)呼叫cuda_tensor.nonzero()依賴於 CUDA tensors 上執行的操作結果的 Python 控制流程，例如
if (cuda_tensor != 0).all()
直接在目標裝置上建立 tensors¶
不要呼叫 torch.rand(size).cuda() 來產生隨機 tensor，而是在目標裝置上直接產生輸出：torch.rand(size, device='cuda')。
這適用於所有建立新 tensor 並接受 device 參數的函式： torch.rand()、torch.zeros()、torch.full() 等類似函式。
使用混合精度和 AMP¶
混合精度利用 Tensor Cores，並在 Volta 及更新的 GPU 架構上提供高達 3 倍的整體加速。 要使用 Tensor Cores，應啟用 AMP，並且矩陣/tensor 維度應滿足呼叫使用 Tensor Cores 的核心的要求。
要使用 Tensor Cores：
在可變輸入長度的情況下預先分配記憶體¶
語音辨識或 NLP 的模型通常使用具有可變序列長度的輸入 tensor 進行訓練。 可變長度對於 PyTorch 緩存分配器來說可能會出現問題，並可能導致效能降低或意外的記憶體不足錯誤。 如果短序列長度的 batch 後面跟著另一個較長序列長度的 batch，則 PyTorch 會被迫釋放先前迭代的中間緩衝區，並重新分配新的緩衝區。 這個過程很耗時，並且會在緩存分配器中造成碎片，這可能會導致記憶體不足錯誤。
一個典型的解決方案是實作預先分配。 它包含以下步驟：
產生一個具有最大序列長度的輸入 batch（通常是隨機的）（對應於訓練資料集中的最大長度或某些預定義的閾值）
使用產生的 batch 執行正向和反向傳遞，不要執行最佳化器或學習速率排程器，此步驟會預先分配最大大小的緩衝區，這些緩衝區可以在後續的訓練迭代中重複使用
清除梯度
繼續進行常規訓練
分散式最佳化¶
使用高效的資料並行後端¶
PyTorch 有兩種實作資料並行訓練的方法：
DistributedDataParallel 提供更好的效能和多 GPU 擴展性。 更多資訊請參考 PyTorch 文件中 CUDA 最佳實踐的相關章節。
如果使用 DistributedDataParallel 和梯度累積進行訓練，則跳過不必要的 all-reduce¶
預設情況下，torch.nn.parallel.DistributedDataParallel 在每次反向傳遞後執行梯度 all-reduce，以計算參與訓練的所有 worker 的平均梯度。 如果訓練使用超過 N 個步驟的梯度累積，則每次訓練步驟後都不需要 all-reduce，只需要在最後一次呼叫反向傳遞之後，在執行最佳化器之前執行 all-reduce。
DistributedDataParallel 提供了 no_sync() 上下文管理器，該管理器會停用特定迭代的梯度 all-reduce。 no_sync() 應應用於梯度累積的前 N-1 次迭代，最後一次迭代應遵循預設執行並執行所需的梯度 all-reduce。
如果使用 DistributedDataParallel(find_unused_parameters=True)，請確保建構函式和執行期間的層順序一致¶
使用 find_unused_parameters=True 的 torch.nn.parallel.DistributedDataParallel 會使用模型建構函式中的層和參數順序，為 DistributedDataParallel 梯度 all-reduce 建立 buckets。DistributedDataParallel 會將 all-reduce 與 backward pass 重疊。只有當給定 bucket 中所有參數的梯度都可用時，才會非同步觸發特定 bucket 的 all-reduce。
為了最大化重疊量，模型建構函式中的順序應大致與執行期間的順序一致。如果順序不一致，則整個 bucket 的 all-reduce 會等待最後到達的梯度，這可能會減少 backward pass 和 all-reduce 之間的重疊，all-reduce 可能最終會暴露出來，從而減慢訓練速度。
使用 find_unused_parameters=False (預設設定) 的 DistributedDataParallel 依賴於基於 backward pass 期間遇到的操作順序的自動 bucket 形成。使用 find_unused_parameters=False，不需要重新排序層或參數以實現最佳效能。
在分散式環境中負載平衡工作負載¶
對於處理序列資料的模型（語音辨識、翻譯、語言模型等），通常可能會發生負載不平衡。如果一個裝置收到一批序列長度比其餘裝置的序列長度更長的資料，則所有裝置都會等待最後完成的 worker。Backward pass 作為具有 DistributedDataParallel 後端的分散式環境中的隱式同步點。
有多種方法可以解決負載平衡問題。核心思想是在每個全域 batch 內，盡可能均勻地將工作負載分配給所有 worker。例如，Transformer 通過形成具有近似恆定 token 數 (以及 batch 中可變數量的序列) 的 batches 來解決不平衡問題，其他模型通過對具有相似序列長度的樣本進行 bucketing，甚至通過按序列長度對資料集進行排序來解決不平衡問題。
腳本總運行時間：（0 分鐘 0.000 秒）
