CUDA 語意¶
torch.cuda 用於設定和執行 CUDA 作業。它會追蹤目前選取的 GPU，並且您配置的所有 CUDA 張量預設都會建立在該裝置上。可以使用 torch.cuda.device 上下文管理器變更選取的裝置。
但是，一旦配置了張量，您就可以對其執行操作，而無需考慮選取的裝置，並且結果將始終放置在與張量相同的裝置上。
預設情況下不允許跨 GPU 操作，但 copy_() 和其他具有類似複製功能的方法（例如 to() 和 cuda()）除外。除非您啟用點對點記憶體存取，否則任何嘗試在分散在不同裝置上的張量上啟動操作的嘗試都將引發錯誤。
您可以在下面找到一個展示此問題的小範例
cuda = torch.device('cuda') # Default CUDA device
cuda0 = torch.device('cuda:0')
cuda2 = torch.device('cuda:2') # GPU 2 (these are 0-indexed)
x = torch.tensor([1., 2.], device=cuda0)
# x.device is device(type='cuda', index=0)
y = torch.tensor([1., 2.]).cuda()
# y.device is device(type='cuda', index=0)
with torch.cuda.device(1):
# allocates a tensor on GPU 1
a = torch.tensor([1., 2.], device=cuda)
# transfers a tensor from CPU to GPU 1
b = torch.tensor([1., 2.]).cuda()
# a.device and b.device are device(type='cuda', index=1)
# You can also use ``Tensor.to`` to transfer a tensor:
b2 = torch.tensor([1., 2.]).to(device=cuda)
# b.device and b2.device are device(type='cuda', index=1)
c = a + b
# c.device is device(type='cuda', index=1)
z = x + y
# z.device is device(type='cuda', index=0)
# even within a context, you can specify the device
# (or give a GPU index to the .cuda call)
d = torch.randn(2, device=cuda2)
e = torch.randn(2).to(cuda2)
f = torch.randn(2).cuda(cuda2)
# d.device, e.device, and f.device are all device(type='cuda', index=2)
Ampere（及更高版本）裝置上的 TensorFloat-32 (TF32)¶
從 PyTorch 1.7 開始，有一個名為 allow_tf32 的新旗標。此旗標在 PyTorch 1.7 到 PyTorch 1.11 中預設為 True，在 PyTorch 1.12 及更高版本中預設為 False。此旗標控制是否允許 PyTorch 使用自 Ampere 以來 NVIDIA GPU 上提供的 TensorFloat32 (TF32) 張量核心來內部計算矩陣乘法（矩陣乘法和批次矩陣乘法）和卷積。
TF32 張量核心旨在透過將輸入資料捨入為具有 10 位尾數並以 FP32 精度累積結果，同時保持 FP32 動態範圍，從而在 torch.float32 張量上的矩陣乘法和卷積上實現更好的效能。
矩陣乘法和卷積分別控制，並且可以透過以下方式存取其相應的旗標
# The flag below controls whether to allow TF32 on matmul. This flag defaults to False
# in PyTorch 1.12 and later.
torch.backends.cuda.matmul.allow_tf32 = True
# The flag below controls whether to allow TF32 on cuDNN. This flag defaults to True.
torch.backends.cudnn.allow_tf32 = True
也可以透過 set_float_32_matmul_precision() 更廣泛地設定矩陣乘法的精度（不僅限於 CUDA）。請注意，除了矩陣乘法和卷積本身之外，內部使用矩陣乘法或卷積的函式和 nn 模組也會受到影響。這些包括 nn.Linear、nn.Conv*、cdist、tensordot、仿射網格和網格樣本、自適應對數 softmax、GRU 和 LSTM。
要瞭解精度和速度，請參閱以下範例程式碼和基準測試資料（在 A100 上）
a_full = torch.randn(10240, 10240, dtype=torch.double, device='cuda')
b_full = torch.randn(10240, 10240, dtype=torch.double, device='cuda')
ab_full = a_full @ b_full
mean = ab_full.abs().mean() # 80.7277
a = a_full.float()
b = b_full.float()
# Do matmul at TF32 mode.
torch.backends.cuda.matmul.allow_tf32 = True
ab_tf32 = a @ b # takes 0.016s on GA100
error = (ab_tf32 - ab_full).abs().max() # 0.1747
relative_error = error / mean # 0.0022
# Do matmul with TF32 disabled.
torch.backends.cuda.matmul.allow_tf32 = False
ab_fp32 = a @ b # takes 0.11s on GA100
error = (ab_fp32 - ab_full).abs().max() # 0.0031
relative_error = error / mean # 0.000039
從上面的範例中，我們可以看到在啟用 TF32 的情況下，A100 上的速度快了約 7 倍，並且與雙精度相比的相對誤差大約大了 2 個數量級。請注意，TF32 與單精度速度的確切比率取決於硬體世代，因為記憶體頻寬與計算的比率以及 TF32 與 FP32 矩陣乘法輸送量的比率可能會因世代或模型而異。如果需要完整的 FP32 精度，則使用者可以透過以下方式停用 TF32
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
要在 C++ 中關閉 TF32 旗標，您可以執行以下操作
at::globalContext().setAllowTF32CuBLAS(false);
at::globalContext().setAllowTF32CuDNN(false);
有關 TF32 的更多資訊，請參閱
FP16 GEMM 中的降低精度減少¶
fp16 GEMM 可能會使用一些中間降低精度減少（例如，以 fp16 而不是 fp32）。這些選擇性的精度降低可以在某些工作負載（特別是具有較大 k 維度的工作負載）和 GPU 架構上以犧牲數值精度和潛在溢出為代價來提高效能。
V100 上的一些範例基準測試資料
[--------------------------- bench_gemm_transformer --------------------------]
[ m , k , n ] | allow_fp16_reduc=True | allow_fp16_reduc=False
1 threads: --------------------------------------------------------------------
[4096, 4048, 4096] | 1634.6 | 1639.8
[4096, 4056, 4096] | 1670.8 | 1661.9
[4096, 4080, 4096] | 1664.2 | 1658.3
[4096, 4096, 4096] | 1639.4 | 1651.0
[4096, 4104, 4096] | 1677.4 | 1674.9
[4096, 4128, 4096] | 1655.7 | 1646.0
[4096, 4144, 4096] | 1796.8 | 2519.6
[4096, 5096, 4096] | 2094.6 | 3190.0
[4096, 5104, 4096] | 2144.0 | 2663.5
[4096, 5112, 4096] | 2149.1 | 2766.9
[4096, 5120, 4096] | 2142.8 | 2631.0
[4096, 9728, 4096] | 3875.1 | 5779.8
[4096, 16384, 4096] | 6182.9 | 9656.5
(times in microseconds).
如果需要全精度減少，則使用者可以使用以下方法停用 fp16 GEMM 中的降低精度減少
torch.backends.cuda.matmul.allow_fp16_reduced_precision_reduction = False
要在 C++ 中切換降低精度減少旗標，可以執行以下操作
at::globalContext().setAllowFP16ReductionCuBLAS(false);
BF16 GEMM 中的降低精度減少¶
BFloat16 GEMM 存在類似的旗標（如上所述）。請注意，對於 BF16，此開關預設設定為 True，如果您在工作負載中觀察到數值不穩定性，則可能需要將其設定為 False。
如果不需要降低精度減少，則使用者可以使用以下方法停用 bf16 GEMM 中的降低精度減少
torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False
要在 C++ 中切換降低精度減少旗標，可以執行以下操作
at::globalContext().setAllowBF16ReductionCuBLAS(true);
非同步執行¶
根據預設，GPU 操作是非同步的。當你呼叫使用 GPU 的函式時，操作會被「排入佇列」到特定裝置，但不一定會立即執行，而是可能會稍後才執行。這允許我們平行執行更多計算，包括在 CPU 或其他 GPU 上的操作。
一般來說，非同步計算的效果對呼叫者來說是不可見的，因為 (1) 每個裝置都按照操作排入佇列的順序執行操作，以及 (2) PyTorch 在 CPU 和 GPU 之間或兩個 GPU 之間複製資料時會自動執行必要的同步。因此，計算會像每個操作都同步執行一樣進行。
你可以透過設定環境變數 CUDA_LAUNCH_BLOCKING=1 來強制同步計算。當 GPU 發生錯誤時，這會很方便。(使用非同步執行時，此類錯誤要等到操作實際執行後才會報告，因此堆疊追蹤不會顯示請求發生的位置。)
非同步計算的結果是，沒有同步的時間測量不準確。為了獲得精確的測量結果，應該在測量之前呼叫 torch.cuda.synchronize()，或使用 torch.cuda.Event 記錄時間，如下所示
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
# Run some things here
end_event.record()
torch.cuda.synchronize() # Wait for the events to be recorded!
elapsed_time_ms = start_event.elapsed_time(end_event)
例外情況是，to() 和 copy_() 等多個函式都允許使用明確的 non_blocking 參數，讓呼叫者在不需要同步時略過同步。另一個例外是 CUDA 流，如下所述。
CUDA 流¶
CUDA 流 是屬於特定裝置的線性執行序列。你通常不需要明確建立一個：根據預設，每個裝置都使用自己的「預設」流。
每個流中的操作會按照建立的順序序列化，但來自不同流的操作可以以任何相對順序同時執行，除非使用明確的同步函式（例如 synchronize() 或 wait_stream()）。例如，以下程式碼不正確
cuda = torch.device('cuda')
s = torch.cuda.Stream() # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
with torch.cuda.stream(s):
# sum() may start execution before normal_() finishes!
B = torch.sum(A)
當「目前流」是預設流時，PyTorch 會在資料移動時自動執行必要的同步，如上所述。但是，當使用非預設流時，使用者有責任確保正確的同步。這個範例的修正版本是
cuda = torch.device('cuda')
s = torch.cuda.Stream() # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
s.wait_stream(torch.cuda.default_stream(cuda)) # NEW!
with torch.cuda.stream(s):
B = torch.sum(A)
A.record_stream(s) # NEW!
有兩個新增項目。torch.cuda.Stream.wait_stream() 呼叫確保在我們開始在側流上執行 sum(A) 之前，normal_() 執行已完成。torch.Tensor.record_stream()（如需更多詳細資訊，請參閱）確保我們在 sum(A) 完成之前不會釋放 A。你也可以稍後使用 torch.cuda.default_stream(cuda).wait_stream(s) 手動等待流（請注意，立即等待是沒有意義的，因為這會阻止流執行與預設流上的其他工作平行執行）。如需何時使用其中一個的更多詳細資訊，請參閱 torch.Tensor.record_stream() 的說明文件。
請注意，即使沒有讀取依賴項，也需要進行這種同步，例如，如下例所示
cuda = torch.device('cuda')
s = torch.cuda.Stream() # Create a new stream.
A = torch.empty((100, 100), device=cuda)
s.wait_stream(torch.cuda.default_stream(cuda)) # STILL REQUIRED!
with torch.cuda.stream(s):
A.normal_(0.0, 1.0)
A.record_stream(s)
儘管對 s 的計算沒有讀取 A 的內容，也沒有其他使用 A 的情況，但仍然需要同步，因為 A 可能對應於 CUDA 快取配置器重新配置的記憶體，並且舊（已釋放）記憶體中有待處理的操作。
反向傳播的流語義¶
每個反向 CUDA 操作都在用於其對應正向操作的相同流上執行。如果你的正向傳播在不同流上平行執行獨立操作，這有助於反向傳播利用相同的平行性。
反向呼叫相對於周圍操作的流語義與任何其他呼叫相同。即使反向操作在多個流上執行（如上一段所述），反向傳播也會插入內部同步以確保這一點。更具體地說，當呼叫 autograd.backward、autograd.grad 或 tensor.backward，並選擇性地提供 CUDA 張量作為初始梯度（例如，autograd.backward(..., grad_tensors=initial_grads)、autograd.grad(..., grad_outputs=initial_grads) 或 tensor.backward(..., gradient=initial_grad)）時，
選擇性地填充初始梯度，
呼叫反向傳播，以及
使用梯度
這些動作與任何一組操作具有相同的流語義關係
s = torch.cuda.Stream()
# Safe, grads are used in the same stream context as backward()
with torch.cuda.stream(s):
loss.backward()
use grads
# Unsafe
with torch.cuda.stream(s):
loss.backward()
use grads
# Safe, with synchronization
with torch.cuda.stream(s):
loss.backward()
torch.cuda.current_stream().wait_stream(s)
use grads
# Safe, populating initial grad and invoking backward are in the same stream context
with torch.cuda.stream(s):
loss.backward(gradient=torch.ones_like(loss))
# Unsafe, populating initial_grad and invoking backward are in different stream contexts,
# without synchronization
initial_grad = torch.ones_like(loss)
with torch.cuda.stream(s):
loss.backward(gradient=initial_grad)
# Safe, with synchronization
initial_grad = torch.ones_like(loss)
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
initial_grad.record_stream(s)
loss.backward(gradient=initial_grad)
BC 備註：在預設流上使用梯度¶
在先前版本的 PyTorch（1.9 及更早版本）中，自動梯度引擎始終將預設流與所有反向操作同步，因此以下模式
with torch.cuda.stream(s):
loss.backward()
use grads
只要 use grads 發生在預設流上就是安全的。在目前的 PyTorch 中，該模式不再安全。如果 backward() 和 use grads 處於不同的流上下文中，則必須同步流
with torch.cuda.stream(s):
loss.backward()
torch.cuda.current_stream().wait_stream(s)
use grads
即使 use grads 位於預設流上也是如此。
記憶體管理¶
PyTorch 使用快取記憶體配置器來加速記憶體配置。這允許在沒有裝置同步的情況下快速釋放記憶體。但是，配置器管理的未使用記憶體在 nvidia-smi 中仍會顯示為已使用。你可以使用 memory_allocated() 和 max_memory_allocated() 監控張量佔用的記憶體，並使用 memory_reserved() 和 max_memory_reserved() 監控快取配置器管理的記憶體總量。呼叫 empty_cache() 會從 PyTorch 釋放所有**未使用**的快取記憶體，以便其他 GPU 應用程式可以使用這些記憶體。但是，張量佔用的 GPU 記憶體不會被釋放，因此它無法增加 PyTorch 可用的 GPU 記憶體量。
為了更好地理解 CUDA 記憶體如何隨著時間推移而被使用，瞭解 CUDA 記憶體使用情況 描述了用於擷取和視覺化記憶體使用軌跡的工具。
對於更進階的使用者，我們透過 memory_stats() 提供更全面的記憶體基準測試。我們還提供了透過 memory_snapshot() 擷取記憶體配置器狀態的完整快照的功能，這可以幫助你理解程式碼產生的底層配置模式。
使用 PYTORCH_CUDA_ALLOC_CONF 優化記憶體使用¶
使用快取配置器可能會干擾記憶體檢查工具，例如 cuda-memcheck。若要使用 cuda-memcheck 除錯記憶體錯誤，請在你的環境中設定 PYTORCH_NO_CUDA_MEMORY_CACHING=1 以停用快取。
快取配置器的行為可以透過環境變數 PYTORCH_CUDA_ALLOC_CONF 來控制。格式為 PYTORCH_CUDA_ALLOC_CONF=<選項>:<值>,<選項2>:<值2>... 可用選項
backend允許選擇底層配置器實現。目前，有效的選項有native（使用 PyTorch 的原生實現）和cudaMallocAsync（使用 CUDA 的內建非同步配置器）。cudaMallocAsync需要 CUDA 11.4 或更新版本。預設值為native。backend適用於處理程序使用的所有裝置，並且不能針對每個裝置指定。max_split_size_mb會阻止原生配置器分割大於此大小（以 MB 為單位）的區塊。這可以減少斷裂，並可能允許一些邊緣工作負載在不耗盡記憶體的情況下完成。效能成本範圍從「零」到「大量」，具體取決於配置模式。預設值是無限制，即所有區塊都可以分割。memory_stats()和memory_summary()方法對於調整很有用。這個選項應該作為最後的手段，用於因「記憶體不足」而中止並顯示大量非活動分割區塊的工作負載。max_split_size_mb僅在backend:native下才有意義。在backend:cudaMallocAsync下，max_split_size_mb會被忽略。roundup_power2_divisions有助於將請求的配置大小四捨五入到最接近的 2 的冪次方，並更好地利用區塊。在原生 CUDACachingAllocator 中，大小以 512 個區塊大小的倍數向上取整，因此這適用於較小的尺寸。但是，這對於附近的大型配置來說效率很低，因為每個配置都會使用不同大小的區塊，並且會盡量減少這些區塊的重複使用。這可能會建立許多未使用的區塊，並浪費 GPU 記憶體容量。此選項可以將配置大小四捨五入到最接近的 2 的冪次方。例如，如果我們需要將大小 1200 向上取整，並且如果分割次數為 4，則大小 1200 介於 1024 和 2048 之間，如果我們在它們之間進行 4 次分割，則這些值為 1024、1280、1536 和 1792。因此，配置大小 1200 將四捨五入到 1280，作為最接近的 2 的冪次方的上限。指定單一值以應用於所有配置大小，或指定一組鍵值對以針對每個 2 的冪次方區間分別設置 2 的冪次方分割。例如，要為 256MB 以下的所有配置設置 1 個分割，為 256MB 到 512MB 之間的配置設置 2 個分割，為 512MB 到 1GB 之間的配置設置 4 個分割，為任何更大的配置設置 8 個分割，請將旋鈕值設置為： [256:1,512:2,1024:4,>:8]。roundup_power2_divisions僅在backend:native下才有意義。在backend:cudaMallocAsync下，roundup_power2_divisions會被忽略。garbage_collection_threshold有助於主動回收未使用的 GPU 記憶體，以避免觸發代價高昂的同步和回收所有操作（release_cached_blocks），這對於延遲關鍵型 GPU 應用程式（例如伺服器）可能是不利的。設置此閾值（例如 0.8）後，如果 GPU 記憶體容量使用量超過閾值（即分配給 GPU 應用程式的總記憶體的 80%），配置器將開始回收 GPU 記憶體區塊。該演算法優先釋放舊的和未使用的區塊，以避免釋放正在被積極重複使用的區塊。閾值應介於大於 0.0 和小於 1.0 之間。garbage_collection_threshold僅在backend:native下才有意義。在backend:cudaMallocAsync下，garbage_collection_threshold會被忽略。expandable_segments（實驗性，預設值：False）如果設置為 True，則此設置會指示配置器建立可以稍後擴展的 CUDA 配置，以便更好地處理作業頻繁更改配置大小的情況，例如批次大小更改。通常，對於大型（>2MB）配置，配置器會呼叫 cudaMalloc 來獲取與使用者請求大小相同的配置。將來，如果這些配置的一部分是空閒的，則可以將其重複用於其他請求。當程式發出許多大小完全相同或大小是該大小的倍數的請求時，這很有效。許多深度學習模型都遵循這種行為。但是，一個常見的例外是當批次大小從一次迭代到下一次迭代略有變化時，例如在批次推斷中。當程式最初以批次大小 N 運行時，它將進行適合該大小的配置。如果將來以大小 N - 1 運行，則現有配置仍然足夠大。但是，如果以大小 N + 1 運行，則必須進行稍微大一點的新配置。並非所有張量的尺寸都相同。有些可能是 (N + 1)*A，而另一些可能是 (N + 1)*A*B，其中 A 和 B 是模型中的一些非批次維度。因為配置器在現有配置足夠大的情況下會重複使用它們，所以一些 (N + 1)*A 配置實際上會放入已經存在的 N*B*A 區段中，儘管並不完美。隨著模型的運行，它將部分填滿所有這些區段，在這些區段的末尾留下無法使用的空閒記憶體切片。配置器在某個時候需要 cudaMalloc 一個新的 (N + 1)*A*B 區段。如果沒有足夠的記憶體，則現在沒有辦法恢復現有區段末尾的空閒記憶體切片。對於深度超過 50 層的模型，此模式可能會重複 50 次以上，從而產生許多碎片。expandable_segments 允許配置器最初建立一個區段，然後在需要更多記憶體時稍後擴展其大小。它不會為每個配置建立一個區段，而是嘗試建立一個根據需要增長的區段（每個流）。現在，當 N + 1 情況運行時，配置將很好地平鋪到一個大區段中，直到它填滿為止。然後請求更多記憶體並附加到區段的末尾。此過程不會建立那麼多無法使用的記憶體碎片，因此更有可能成功找到此記憶體。
pinned_use_cuda_host_register 選項是一個布林值標記，用於確定是否使用 CUDA API 的 cudaHostRegister 函數來配置固定記憶體，而不是預設的 cudaHostAlloc。如果設置為 True，則使用常規 malloc 配置記憶體，然後在呼叫 cudaHostRegister 之前將頁面映射到記憶體。這種頁面的預先映射有助於減少執行 cudaHostRegister 期間的鎖定時間。
pinned_num_register_threads 選項僅在 pinned_use_cuda_host_register 設置為 True 時才有效。預設情況下，使用一個執行緒來映射頁面。此選項允許使用更多執行緒來並行化頁面映射操作，以減少固定記憶體的總體配置時間。根據基準測試結果，此選項的良好值為 8。
備註
CUDA 記憶體管理 API 報告的一些統計資訊特定於 backend:native，在 backend:cudaMallocAsync 下沒有意義。有關詳細資訊，請參閱每個函數的文檔字串。
為 CUDA 使用自定義記憶體配置器¶
可以將配置器定義為 C/C++ 中的簡單函數，並將其編譯為共享庫，下面的程式碼顯示了一個基本的配置器，它只跟踪所有記憶體操作。
#include <sys/types.h>
#include <cuda_runtime_api.h>
#include <iostream>
// Compile with g++ alloc.cc -o alloc.so -I/usr/local/cuda/include -shared -fPIC
extern "C" {
void* my_malloc(ssize_t size, int device, cudaStream_t stream) {
void *ptr;
cudaMalloc(&ptr, size);
std::cout<<"alloc "<<ptr<<size<<std::endl;
return ptr;
}
void my_free(void* ptr, ssize_t size, int device, cudaStream_t stream) {
std::cout<<"free "<<ptr<< " "<<stream<<std::endl;
cudaFree(ptr);
}
}
這可以在 Python 中透過 torch.cuda.memory.CUDAPluggableAllocator 使用。使用者負責提供 .so 文件的路徑以及與上面指定的簽章匹配的 alloc/free 函數的名稱。
import torch
# Load the allocator
new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
'alloc.so', 'my_malloc', 'my_free')
# Swap the current allocator
torch.cuda.memory.change_current_allocator(new_alloc)
# This will allocate memory in the device using the new allocator
b = torch.zeros(10, device='cuda')
import torch
# Do an initial memory allocator
b = torch.zeros(10, device='cuda')
# Load the allocator
new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
'alloc.so', 'my_malloc', 'my_free')
# This will error since the current allocator was already instantiated
torch.cuda.memory.change_current_allocator(new_alloc)
cuBLAS 工作區¶
對於 cuBLAS 控制代碼和 CUDA 流的每種組合，如果該控制代碼和流組合執行需要工作區的 cuBLAS 核心，則會配置一個 cuBLAS 工作區。為了避免重複配置工作區，除非呼叫 torch._C._cuda_clearCublasWorkspaces()，否則不會釋放這些工作區。每個配置的工作區大小可以透過環境變數 CUBLAS_WORKSPACE_CONFIG 指定，格式為 :[SIZE]:[COUNT]。例如，每個配置的預設工作區大小為 CUBLAS_WORKSPACE_CONFIG=:4096:2:16:8，它指定了總大小為 2 * 4096 + 8 * 16 KiB。若要強制 cuBLAS 避免使用工作區，請設置 CUBLAS_WORKSPACE_CONFIG=:0:0。
cuFFT 計劃快取¶
對於每個 CUDA 裝置，都會使用一個 cuFFT 計劃的 LRU 快取來加速在具有相同配置的相同幾何形狀的 CUDA 張量上重複運行 FFT 方法（例如，torch.fft.fft()）。因為某些 cuFFT 計劃可能會配置 GPU 記憶體，所以這些快取具有最大容量。
您可以使用以下 API 控制和查詢當前裝置快取的屬性
torch.backends.cuda.cufft_plan_cache.max_size給出快取的容量（在 CUDA 10 及更新版本上預設為 4096，在舊版 CUDA 上預設為 1023）。直接設置此值會修改容量。torch.backends.cuda.cufft_plan_cache.size給出當前駐留在快取中的計劃數量。torch.backends.cuda.cufft_plan_cache.clear()清除快取。
若要控制和查詢非預設裝置的計劃快取，您可以使用 torch.device 物件或裝置索引為 torch.backends.cuda.cufft_plan_cache 物件編制索引，並存取上述屬性之一。例如，若要設置裝置 1 的快取容量，可以編寫 torch.backends.cuda.cufft_plan_cache[1].max_size = 10。
即時編譯¶
當在 CUDA 張量上執行時，PyTorch 會即時編譯一些操作，例如 torch.special.zeta。此編譯可能非常耗時（根據您的硬體和軟體，最長可能需要幾秒鐘），並且對於單個運算符可能會發生多次，因為許多 PyTorch 運算符實際上是從各種核心（每個核心都必須編譯一次，具體取決於其輸入）中進行選擇的。此編譯每個進程發生一次，或者如果使用核心快取，則僅發生一次。
預設情況下，如果定義了 XDG_CACHE_HOME，則 PyTorch 會在 $XDG_CACHE_HOME/torch/kernels 中建立核心快取；如果未定義 XDG_CACHE_HOME，則在 $HOME/.cache/torch/kernels 中建立核心快取（Windows 除外，Windows 尚不支援核心快取）。可以使用兩個環境變數直接控制快取行為。如果 USE_PYTORCH_KERNEL_CACHE 設置為 0，則不會使用快取，如果設置了 PYTORCH_KERNEL_CACHE_PATH，則該路徑將用作核心快取，而不是預設位置。
最佳實務¶
與裝置無關的程式碼¶
由於 PyTorch 的結構，您可能需要明確地編寫與裝置無關（CPU 或 GPU）的程式碼；例如，建立一個新的張量作為遞迴神經網路的初始隱藏狀態。
第一步是確定是否應該使用 GPU。一種常見的模式是使用 Python 的 argparse 模組讀取使用者參數，並設定一個可以用於停用 CUDA 的旗標，並搭配使用 is_available()。在以下範例中，args.device 會產生一個 torch.device 物件，可用於將張量移動到 CPU 或 CUDA。
import argparse
import torch
parser = argparse.ArgumentParser(description='PyTorch Example')
parser.add_argument('--disable-cuda', action='store_true',
help='Disable CUDA')
args = parser.parse_args()
args.device = None
if not args.disable_cuda and torch.cuda.is_available():
args.device = torch.device('cuda')
else:
args.device = torch.device('cpu')
備註
在評估 CUDA 在給定環境中的可用性時 (is_available())，PyTorch 的預設行為是呼叫 CUDA Runtime API 方法 cudaGetDeviceCount。因為如果尚未初始化，此呼叫會初始化 CUDA Driver API（透過 cuInit），因此後續已執行 is_available() 的程序分支將因 CUDA 初始化錯誤而失敗。
您可以在匯入執行 is_available() 的 PyTorch 模組之前（或直接執行之前），在您的環境中設定 PYTORCH_NVML_BASED_CUDA_CHECK=1，以指示 is_available() 嘗試進行基於 NVML 的評估 (nvmlDeviceGetCount_v2)。如果基於 NVML 的評估成功（即 NVML 探索/初始化未失敗），is_available() 呼叫將不會影響後續的分支。
如果 NVML 探索/初始化失敗，is_available() 將回退到標準 CUDA Runtime API 評估，並且上述的分支限制將適用。
請注意，上述基於 NVML 的 CUDA 可用性評估提供的保證比預設的 CUDA Runtime API 方法（需要 CUDA 初始化成功）弱。在某些情況下，基於 NVML 的檢查可能會成功，而後來的 CUDA 初始化會失敗。
現在我們有了 args.device，我們可以使用它在所需的裝置上建立張量。
x = torch.empty((8, 42), device=args.device)
net = Network().to(device=args.device)
這可以在許多情況下用於產生與裝置無關的程式碼。以下是在使用資料載入器時的範例
cuda0 = torch.device('cuda:0') # CUDA GPU 0
for i, x in enumerate(train_loader):
x = x.to(cuda0)
在系統上使用多個 GPU 時，您可以使用 CUDA_VISIBLE_DEVICES 環境旗標來管理哪些 GPU 可供 PyTorch 使用。如上所述，若要手動控制在哪个 GPU 上建立張量，最佳做法是使用 torch.cuda.device 上下文管理器。
print("Outside device is 0") # On device 0 (default in most scenarios)
with torch.cuda.device(1):
print("Inside device is 1") # On device 1
print("Outside device is still 0") # On device 0
如果您有一個張量，並且想要在同一個裝置上建立一個相同類型的張量，則可以使用 torch.Tensor.new_* 方法（請參閱 torch.Tensor）。雖然前面提到的 torch.* 工廠函數（建立運算）取決於目前的 GPU 上下文和您傳入的屬性參數，但 torch.Tensor.new_* 方法會保留張量的裝置和其他屬性。
在建立模組時，如果需要在正向傳遞過程中內部建立新的張量，建議使用這種做法。
cuda = torch.device('cuda')
x_cpu = torch.empty(2)
x_gpu = torch.empty(2, device=cuda)
x_cpu_long = torch.empty(2, dtype=torch.int64)
y_cpu = x_cpu.new_full([3, 2], fill_value=0.3)
print(y_cpu)
tensor([[ 0.3000, 0.3000],
[ 0.3000, 0.3000],
[ 0.3000, 0.3000]])
y_gpu = x_gpu.new_full([3, 2], fill_value=-5)
print(y_gpu)
tensor([[-5.0000, -5.0000],
[-5.0000, -5.0000],
[-5.0000, -5.0000]], device='cuda:0')
y_cpu_long = x_cpu_long.new_tensor([[1, 2, 3]])
print(y_cpu_long)
tensor([[ 1, 2, 3]])
如果您想建立一個與另一個張量類型和大小相同的張量，並用 1 或 0 填充它，則可以使用 ones_like() 或 zeros_like() 作為方便的輔助函數（它們也保留 torch.device 和 torch.dtype）。
x_cpu = torch.empty(2, 3)
x_gpu = torch.empty(2, 3)
y_cpu = torch.ones_like(x_cpu)
y_gpu = torch.zeros_like(x_gpu)
使用固定記憶體緩衝區¶
警告
這是一個進階技巧。如果您過度使用固定記憶體，在 RAM 不足時可能會導致嚴重的問題，並且您應該注意固定記憶體通常是一項昂貴的操作。
當主機到 GPU 的複製源自固定（頁面鎖定）記憶體時，速度會快得多。CPU 張量和儲存體公開了一個 pin_memory() 方法，該方法會傳回物件的副本，並將資料放入固定的區域中。
此外，一旦您固定了張量或儲存體，就可以使用非同步 GPU 複製。只需將額外的 non_blocking=True 參數傳遞給 to() 或 cuda() 呼叫即可。這可以用於將資料傳輸與計算重疊。
您可以透過將 pin_memory=True 傳遞給 DataLoader 的建構函數，使其傳回放置在固定記憶體中的批次。
使用 nn.parallel.DistributedDataParallel 而不是 multiprocessing 或 nn.DataParallel¶
大多數涉及批次輸入和多個 GPU 的用例，預設情況下應該使用 DistributedDataParallel 來利用多個 GPU。
使用 multiprocessing 搭配 CUDA 模型有一些需要注意的事項；除非仔細滿足資料處理需求，否則您的程式很可能會有不正確或未定義的行為。
建議使用 DistributedDataParallel 而不是 DataParallel 進行多 GPU 訓練，即使只有一個節點也是如此。
DistributedDataParallel 和 DataParallel 的區別在於：DistributedDataParallel 使用多程序，其中為每個 GPU 建立一個程序，而 DataParallel 使用多執行緒。透過使用多程序，每個 GPU 都有其專屬的程序，這避免了 Python 解釋器的 GIL 所造成的效能損失。
如果您使用 DistributedDataParallel，您可以使用 torch.distributed.launch 工具來啟動您的程式，請參閱 第三方後端。
CUDA 圖形¶
CUDA 圖形是 CUDA 串流及其相依串流所執行工作的記錄（主要是核心及其參數）。有關底層 CUDA API 的一般原則和詳細資訊，請參閱 CUDA 圖形入門 以及 CUDA C 程式設計指南的 圖形章節。
PyTorch 支援使用 串流擷取 來建構 CUDA 圖形，這會將 CUDA 串流置於 *擷取模式*。發送到擷取串流的 CUDA 工作實際上並未在 GPU 上執行。相反的，工作會記錄在圖形中。
擷取後，圖形可以 *啟動* 以根據需要多次執行 GPU 工作。每次重播都會使用相同的參數執行相同的核心。對於指標參數，這表示使用相同的記憶體位址。透過在每次重播之前用新資料（例如，來自新批次的資料）填充輸入記憶體，您可以在新資料上重新執行相同的工作。
為什麼要使用 CUDA 圖形？¶
重放圖形犧牲了典型 Eager Execution 的動態靈活性，以換取**大幅降低 CPU 開銷**。圖形的引數和核心是固定的，因此圖形重放會跳過所有層級的引數設定和核心調度，包括 Python、C++ 和 CUDA 驅動程式的開銷。在底層，重放會使用單一呼叫 cudaGraphLaunch 將整個圖形的工作提交給 GPU。核心在重放時在 GPU 上的執行速度也會略快，但消除 CPU 開銷是主要優勢。
如果您的網路的全部或部分是圖形安全的（通常這意味著靜態形狀和靜態控制流程，但請參閱其他限制），並且您懷疑其運行時至少在一定程度上受限於 CPU，則您應該嘗試使用 CUDA 圖形。
PyTorch API¶
警告
此 API 處於測試階段，可能會在未來版本中變更。
PyTorch 透過原始 torch.cuda.CUDAGraph 類別和兩個便利的包裝器，torch.cuda.graph 和 torch.cuda.make_graphed_callables 公開圖形。
torch.cuda.graph 是一個簡單、通用的上下文管理器，可以在其上下文中捕獲 CUDA 工作。在捕獲之前，請透過運行幾個 Eager 迭代來預熱要捕獲的工作負載。預熱必須在側流上進行。由於圖形在每次重放時都會從相同的記憶體位址讀取和寫入，因此您必須在捕獲期間保持對保存輸入和輸出數據的張量的長期引用。若要對新的輸入數據運行圖形，請將新的數據複製到捕獲的輸入張量，重放圖形，然後從捕獲的輸出張量讀取新的輸出。範例
g = torch.cuda.CUDAGraph()
# Placeholder input used for capture
static_input = torch.empty((5,), device="cuda")
# Warmup before capture
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
for _ in range(3):
static_output = static_input * 2
torch.cuda.current_stream().wait_stream(s)
# Captures the graph
# To allow capture, automatically sets a side stream as the current stream in the context
with torch.cuda.graph(g):
static_output = static_input * 2
# Fills the graph's input memory with new data to compute on
static_input.copy_(torch.full((5,), 3, device="cuda"))
g.replay()
# static_output holds the results
print(static_output) # full of 3 * 2 = 6
# Fills the graph's input memory with more data to compute on
static_input.copy_(torch.full((5,), 4, device="cuda"))
g.replay()
print(static_output) # full of 4 * 2 = 8
有關真實和進階模式，請參閱 全網路捕獲、與 torch.cuda.amp 一起使用 和 與多個流一起使用。
make_graphed_callables 更為複雜。make_graphed_callables 接受 Python 函數和 torch.nn.Module。對於每個傳入的函數或模組，它會建立前向傳遞和後向傳遞工作的獨立圖形。請參閱 部分網路捕獲。
限制¶
如果一組操作不違反以下任何限制，則該組操作是*可捕獲的*。
限制適用於 torch.cuda.graph 上下文中的所有工作，以及您傳遞給 torch.cuda.make_graphed_callables() 的任何可呼叫物件的前向和後向傳遞中的所有工作。
違反任何這些限制都可能會導致執行階段錯誤
捕獲必須在非預設流上進行。（如果您使用原始的
CUDAGraph.capture_begin和CUDAGraph.capture_end呼叫，則只需擔心這一點。graph和make_graphed_callables()會為您設定側流。）禁止使用與 GPU 同步 CPU 的操作（例如，
.item()呼叫）。允許使用 CUDA RNG 操作，並且在圖形中使用多個
torch.Generator執行個體時，必須在圖形捕獲之前使用CUDAGraph.register_generator_state註冊它們。避免在捕獲期間使用Generator.get_state和Generator.set_state；相反，請使用Generator.graphsafe_set_state和Generator.graphsafe_get_state來安全地管理圖形上下文中的產生器狀態。這可確保在 CUDA 圖形中正確執行 RNG 操作和產生器管理。
違反任何這些限制都可能會導致靜默的數值錯誤或未定義的行為
在一個行程中，一次只能進行一個捕獲。
在捕獲進行時，此行程（在任何執行緒上）都不能執行未捕獲的 CUDA 工作。
CPU 工作不會被捕獲。如果捕獲的操作包含 CPU 工作，則在重放期間將省略該工作。
每次重放都會從相同的（虛擬）記憶體位址讀取和寫入。
禁止使用動態控制流程（基於 CPU 或 GPU 數據）。
禁止使用動態形狀。圖形假設捕獲的操作序列中的每個張量在每次重放時都具有相同的大小和佈局。
允許在捕獲中使用多個流，但有一些限制。
非限制¶
捕獲後，圖形可以在任何流上重放。
全網路捕獲¶
如果您的整個網路都是可捕獲的，則您可以捕獲和重放整個迭代
N, D_in, H, D_out = 640, 4096, 2048, 1024
model = torch.nn.Sequential(torch.nn.Linear(D_in, H),
torch.nn.Dropout(p=0.2),
torch.nn.Linear(H, D_out),
torch.nn.Dropout(p=0.1)).cuda()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
# Placeholders used for capture
static_input = torch.randn(N, D_in, device='cuda')
static_target = torch.randn(N, D_out, device='cuda')
# warmup
# Uses static_input and static_target here for convenience,
# but in a real setting, because the warmup includes optimizer.step()
# you must use a few batches of real data.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
for i in range(3):
optimizer.zero_grad(set_to_none=True)
y_pred = model(static_input)
loss = loss_fn(y_pred, static_target)
loss.backward()
optimizer.step()
torch.cuda.current_stream().wait_stream(s)
# capture
g = torch.cuda.CUDAGraph()
# Sets grads to None before capture, so backward() will create
# .grad attributes with allocations from the graph's private pool
optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
static_y_pred = model(static_input)
static_loss = loss_fn(static_y_pred, static_target)
static_loss.backward()
optimizer.step()
real_inputs = [torch.rand_like(static_input) for _ in range(10)]
real_targets = [torch.rand_like(static_target) for _ in range(10)]
for data, target in zip(real_inputs, real_targets):
# Fills the graph's input memory with new data to compute on
static_input.copy_(data)
static_target.copy_(target)
# replay() includes forward, backward, and step.
# You don't even need to call optimizer.zero_grad() between iterations
# because the captured backward refills static .grad tensors in place.
g.replay()
# Params have been updated. static_y_pred, static_loss, and .grad
# attributes hold values from computing on this iteration's data.
部分網路捕獲¶
如果您的網路的某些部分不安全而無法捕獲（例如，由於動態控制流程、動態形狀、CPU 同步或必要的 CPU 端邏輯），則您可以急切地運行不安全的部份，並使用 torch.cuda.make_graphed_callables() 僅將可捕獲的部份繪製為圖形。
根據預設，make_graphed_callables() 返回的可呼叫物件是可感知自動梯度的，並且可以在訓練迴圈中用作您傳遞的函數或 nn.Module 的直接替代品。
make_graphed_callables() 在內部建立 CUDAGraph 物件、運行預熱迭代，並根據需要維護靜態輸入和輸出。因此（與 torch.cuda.graph 不同），您不需要手動處理這些。
在以下範例中，依賴於數據的動態控制流程意味著網路無法端到端捕獲，但 make_graphed_callables() 允許我們捕獲圖形安全的區段並將其作為圖形運行，而無需理會
N, D_in, H, D_out = 640, 4096, 2048, 1024
module1 = torch.nn.Linear(D_in, H).cuda()
module2 = torch.nn.Linear(H, D_out).cuda()
module3 = torch.nn.Linear(H, D_out).cuda()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(chain(module1.parameters(),
module2.parameters(),
module3.parameters()),
lr=0.1)
# Sample inputs used for capture
# requires_grad state of sample inputs must match
# requires_grad state of real inputs each callable will see.
x = torch.randn(N, D_in, device='cuda')
h = torch.randn(N, H, device='cuda', requires_grad=True)
module1 = torch.cuda.make_graphed_callables(module1, (x,))
module2 = torch.cuda.make_graphed_callables(module2, (h,))
module3 = torch.cuda.make_graphed_callables(module3, (h,))
real_inputs = [torch.rand_like(x) for _ in range(10)]
real_targets = [torch.randn(N, D_out, device="cuda") for _ in range(10)]
for data, target in zip(real_inputs, real_targets):
optimizer.zero_grad(set_to_none=True)
tmp = module1(data) # forward ops run as a graph
if tmp.sum().item() > 0:
tmp = module2(tmp) # forward ops run as a graph
else:
tmp = module3(tmp) # forward ops run as a graph
loss = loss_fn(tmp, target)
# module2's or module3's (whichever was chosen) backward ops,
# as well as module1's backward ops, run as graphs
loss.backward()
optimizer.step()
與 torch.cuda.amp 一起使用¶
對於典型的優化器，GradScaler.step 會將 CPU 與 GPU 同步，這在捕獲期間是被禁止的。為避免錯誤，請使用 部分網路捕獲，或者（如果前向、損失和後向都是可捕獲的）捕獲前向、損失和後向，但不要捕獲優化器步驟
# warmup
# In a real setting, use a few batches of real data.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
for i in range(3):
optimizer.zero_grad(set_to_none=True)
with torch.cuda.amp.autocast():
y_pred = model(static_input)
loss = loss_fn(y_pred, static_target)
scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
torch.cuda.current_stream().wait_stream(s)
# capture
g = torch.cuda.CUDAGraph()
optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
with torch.cuda.amp.autocast():
static_y_pred = model(static_input)
static_loss = loss_fn(static_y_pred, static_target)
scaler.scale(static_loss).backward()
# don't capture scaler.step(optimizer) or scaler.update()
real_inputs = [torch.rand_like(static_input) for _ in range(10)]
real_targets = [torch.rand_like(static_target) for _ in range(10)]
for data, target in zip(real_inputs, real_targets):
static_input.copy_(data)
static_target.copy_(target)
g.replay()
# Runs scaler.step and scaler.update eagerly
scaler.step(optimizer)
scaler.update()
與多個流一起使用¶
捕獲模式會自動傳播到與捕獲流同步的任何流。在捕獲中，您可以透過發出對不同流的呼叫來公開平行性，但整體流依賴性 DAG 必須在捕獲開始後從初始捕獲流分支出來，並在捕獲結束前重新加入初始流
with torch.cuda.graph(g):
# at context manager entrance, torch.cuda.current_stream()
# is the initial capturing stream
# INCORRECT (does not branch out from or rejoin initial stream)
with torch.cuda.stream(s):
cuda_work()
# CORRECT:
# branches out from initial stream
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
cuda_work()
# rejoins initial stream before capture ends
torch.cuda.current_stream().wait_stream(s)
備註
為避免在 nsight systems 或 nvprof 中查看重放的高級用戶感到困惑：與 Eager Execution 不同，圖形將捕獲中的非平凡流 DAG 解釋為提示，而不是命令。在重放期間，圖形可能會將獨立的操作重新組織到不同的流上，或者以不同的順序（同時遵守您原始 DAG 的整體依賴性）將它們排入佇列。
與 DistributedDataParallel 一起使用¶
NCCL < 2.9.6¶
2.9.6 之前的 NCCL 版本不允許捕獲集體運算。您必須使用 部分網路捕獲，它會延遲 allreduce 在後向圖形化區段之外發生。
在使用 DDP 包裝網路*之前*，在可繪製圖形的網路區段上呼叫 make_graphed_callables()。
NCCL >= 2.9.6¶
2.9.6 或更高版本的 NCCL 允許在圖形中使用集體運算。捕獲整個後向傳遞的方法是一種可行的選擇，但需要三個設定步驟。
停用 DDP 的內部非同步錯誤處理
os.environ["NCCL_ASYNC_ERROR_HANDLING"] = "0" torch.distributed.init_process_group(...)
在完整後向捕獲之前，必須在側流上下文中建構 DDP
with torch.cuda.stream(s): model = DistributedDataParallel(model)
在捕獲之前，您的預熱必須至少運行 11 個啟用 DDP 的 Eager 迭代。
圖形記憶體管理¶
擷取的圖形每次重播時都會作用於相同的虛擬地址。如果 PyTorch 釋放記憶體，則後續重播可能會遇到非法記憶體存取。如果 PyTorch 將記憶體重新分配給新的張量，則重播可能會損壞這些張量所看到的值。因此，圖形使用的虛擬地址必須在重播時為圖形保留。PyTorch 快取分配器會偵測何時正在進行擷取，並從圖形專用記憶體池滿足擷取的分配，從而實現這一點。專用池會一直存在，直到其 CUDAGraph 物件和在擷取期間建立的所有張量都超出範圍。
專用池會自動維護。根據預設，分配器會為每個擷取建立一個單獨的專用池。如果您擷取多個圖形，這種保守的方法可確保圖形重播永遠不會損壞彼此的值，但有時會不必要地浪費記憶體。
跨擷取共用記憶體¶
為了節省存放在專用池中的記憶體，torch.cuda.graph 和 torch.cuda.make_graphed_callables() 可選擇允許不同的擷取共用同一個專用池。如果您知道一組圖形將始終按照擷取的順序重播，並且永遠不會同時重播，則它們共用一個專用池是安全的。
torch.cuda.graph 的 pool 參數是用於使用特定專用池的提示，並且可用於跨圖形共用記憶體，如下所示
g1 = torch.cuda.CUDAGraph()
g2 = torch.cuda.CUDAGraph()
# (create static inputs for g1 and g2, run warmups of their workloads...)
# Captures g1
with torch.cuda.graph(g1):
static_out_1 = g1_workload(static_in_1)
# Captures g2, hinting that g2 may share a memory pool with g1
with torch.cuda.graph(g2, pool=g1.pool()):
static_out_2 = g2_workload(static_in_2)
static_in_1.copy_(real_data_1)
static_in_2.copy_(real_data_2)
g1.replay()
g2.replay()
使用 torch.cuda.make_graphed_callables() 時，如果您想繪製多個可呼叫物件的圖形，並且您知道它們將始終以相同的順序運行（並且永遠不會同時運行），請按照它們在實際工作負載中的運行順序將它們作為元組傳遞，並且 make_graphed_callables() 將使用共用的專用池擷取它們的圖形。
如果在實際工作負載中，您的可呼叫物件的運行順序偶爾會發生變化，或者它們將同時運行，則不允許將它們作為元組傳遞給 make_graphed_callables() 的單次調用。相反，您必須為每個可呼叫物件單獨調用 make_graphed_callables()。
