torch.cuda

此包增加了對 CUDA 張量型別的支援。

它實現了與 CPU 張量相同的功能，但利用 GPU 進行計算。

它是延遲初始化的，所以你可以隨時匯入它，並使用 is_available() 來判斷你的系統是否支援 CUDA。

CUDA 語義 提供了更多關於使用 CUDA 的詳細資訊。

StreamContext	選擇給定流的上下文管理器。
can_device_access_peer	檢查兩個裝置之間是否可以進行對等訪問。
current_blas_handle	返回指向當前 cuBLAS 控制代碼的 cublasHandle_t 指標
current_device	返回當前選定裝置的索引。
current_stream	返回給定裝置當前選定的 Stream。
cudart	檢索 CUDA 執行時 API 模組。
default_stream	返回給定裝置的預設 Stream。
device	更改選定裝置的上下文管理器。
device_count	返回可用 GPU 的數量。
device_memory_used	返回由 nvidia-smi 或 amd-smi 報告的已使用的全域性 (裝置) 記憶體位元組數。
device_of	將當前裝置更改為給定物件裝置的上下文管理器。
get_arch_list	返回此庫編譯所支援的 CUDA 架構列表。
get_device_capability	獲取裝置的 CUDA 能力。
get_device_name	獲取裝置的名稱。
get_device_properties	獲取裝置的屬性。
get_gencode_flags	返回此庫編譯時使用的 NVCC gencode 標誌。
get_stream_from_external	從外部分配的 CUDA 流返回一個 Stream。
get_sync_debug_mode	返回 CUDA 同步操作的當前除錯模式值。
init	初始化 PyTorch 的 CUDA 狀態。
ipc_collect	強制收集被 CUDA IPC 釋放後的 GPU 記憶體。
is_available	返回一個布林值，指示 CUDA 當前是否可用。
is_initialized	返回 PyTorch 的 CUDA 狀態是否已初始化。
is_tf32_supported	返回一個布林值，指示當前的 CUDA/ROCm 裝置是否支援 dtype tf32。
memory_usage	返回在過去取樣週期內，全域性 (裝置) 記憶體被讀取或寫入的時間百分比，由 nvidia-smi 報告。
set_device	設定當前裝置。
set_stream	設定當前流。這是一個設定流的包裝 API。
set_sync_debug_mode	設定 CUDA 同步操作的除錯模式。
stream	包裝 StreamContext 上下文管理器，用於選擇給定流。
synchronize	等待 CUDA 裝置上所有流中的所有核心完成。
utilization	返回在過去取樣週期內，一個或多個核心在 GPU 上執行的時間百分比，由 nvidia-smi 報告。
temperature	返回 GPU 感測器以攝氏度 (Degrees C) 為單位的平均溫度。
power_draw	返回 GPU 感測器以毫瓦 (MilliWatts) 為單位的平均功耗。
clock_rate	返回在過去取樣週期內，GPU SM 以兆赫茲 (megahertz) 為單位的時鐘速度，由 nvidia-smi 報告。
OutOfMemoryError	裝置記憶體不足時引發的異常

隨機數生成器

get_rng_state	返回指定 GPU 的隨機數生成器狀態，作為 ByteTensor。
get_rng_state_all	返回一個 ByteTensor 列表，代表所有裝置的隨機數狀態。
set_rng_state	設定指定 GPU 的隨機數生成器狀態。
set_rng_state_all	設定所有裝置的隨機數生成器狀態。
manual_seed	為當前 GPU 設定用於生成隨機數的種子。
manual_seed_all	為所有 GPU 設定用於生成隨機數的種子。
seed	為當前 GPU 將隨機數生成種子設定為一個隨機數。
seed_all	為所有 GPU 將隨機數生成種子設定為一個隨機數。
initial_seed	返回當前 GPU 的當前隨機種子。

通訊集合操作

comm.broadcast	將張量廣播到指定的 GPU 裝置。
comm.broadcast_coalesced	將一系列張量廣播到指定的 GPU。
comm.reduce_add	對來自多個 GPU 的張量求和。
comm.scatter	將張量分散到多個 GPU。
comm.gather	從多個 GPU 裝置收集張量。

流和事件

Stream	CUDA 流的包裝器。
ExternalStream	外部分配的 CUDA 流的包裝器。
Event	CUDA 事件的包裝器。

圖 (beta)

is_current_stream_capturing	如果當前 CUDA 流正在進行 CUDA 圖捕獲，則返回 True，否則返回 False。
graph_pool_handle	返回一個不透明的令牌，表示圖記憶體池的 ID。
CUDAGraph	CUDA 圖的包裝器。
graph	一個上下文管理器，用於將 CUDA 工作捕獲到 torch.cuda.CUDAGraph 物件中，以便後續回放。
make_graphed_callables	接受可呼叫物件（函式或 nn.Module）並返回圖化版本。

記憶體管理

empty_cache	釋放當前由快取分配器持有的所有未佔用的快取記憶體，以便這些記憶體可用於其他 GPU 應用並在 nvidia-smi 中可見。
get_per_process_memory_fraction	獲取程序的記憶體比例。
list_gpu_processes	返回給定裝置的執行程序及其 GPU 記憶體使用的可讀性報告。
mem_get_info	使用 cudaMemGetInfo 返回給定裝置的全域性空閒和總 GPU 記憶體。
memory_stats	返回給定裝置的 CUDA 記憶體分配器統計資訊字典。
host_memory_stats	返回給定裝置的 CUDA 記憶體分配器統計資訊字典。
memory_summary	返回給定裝置的當前記憶體分配器統計資訊的可讀性報告。
memory_snapshot	返回所有裝置的 CUDA 記憶體分配器狀態快照。
memory_allocated	返回給定裝置當前由張量佔用的 GPU 記憶體位元組數。
max_memory_allocated	返回給定裝置由張量佔用的最大 GPU 記憶體位元組數。
reset_max_memory_allocated	重置給定裝置跟蹤由張量佔用的最大 GPU 記憶體的起始點。
memory_reserved	返回給定裝置當前由快取分配器管理的 GPU 記憶體位元組數。
max_memory_reserved	返回給定裝置由快取分配器管理的最大 GPU 記憶體位元組數。
set_per_process_memory_fraction	設定程序的記憶體比例。
memory_cached	已棄用；請參閱 memory_reserved()。
max_memory_cached	已棄用；請參閱 max_memory_reserved()。
reset_max_memory_cached	重置給定裝置跟蹤由快取分配器管理的最大 GPU 記憶體的起始點。
reset_peak_memory_stats	重置 CUDA 記憶體分配器跟蹤的“峰值”統計資訊。
reset_peak_host_memory_stats	重置主機記憶體分配器跟蹤的“峰值”統計資訊。
caching_allocator_alloc	使用 CUDA 記憶體分配器執行記憶體分配。
caching_allocator_delete	刪除使用 CUDA 記憶體分配器分配的記憶體。
get_allocator_backend	返回一個字串，描述由 PYTORCH_CUDA_ALLOC_CONF 設定的活動分配器後端。
CUDAPluggableAllocator	從 .so 檔案載入的 CUDA 記憶體分配器。
change_current_allocator	將當前使用的記憶體分配器更改為提供的分配器。
MemPool	MemPool 表示快取分配器中的一個記憶體池。
MemPoolContext	MemPoolContext 儲存當前活動的記憶體池並暫存之前的記憶體池。

caching_allocator_enable

啟用或停用 CUDA 記憶體分配器。

class torch.cuda.use_mem_pool(pool, device=None)

一個將分配路由到給定記憶體池的上下文管理器。

引數

• pool (torch.cuda.MemPool) – 一個要被啟用的 MemPool 物件，以便將分配路由到此記憶體池。

• device (torch.device 或 int, 可選) – 選定裝置。如果 device 為 None (預設)，則使用 current_device() 給定的當前裝置的 MemPool。

NVIDIA 工具擴充套件 (NVTX)

nvtx.mark	描述在某個時間點發生的瞬時事件。
nvtx.range_push	將一個範圍推入巢狀範圍跨度棧。
nvtx.range_pop	從巢狀範圍跨度棧彈出一個範圍。
nvtx.range	上下文管理器 / 裝飾器，在其作用域開始時推入一個 NVTX 範圍，並在結束時彈出。

Jiterator (beta)

jiterator._create_jit_fn	為逐元素操作建立一個由 jiterator 生成的 cuda kernel。
jiterator._create_multi_output_jit_fn	為支援返回一個或多個輸出的逐元素操作建立一個由 jiterator 生成的 cuda kernel。

TunableOp

某些操作可以使用不止一個庫或不止一種技術來實現。例如，GEMM 可以針對 CUDA 或 ROCm 使用 cublas/cublasLt 庫或 hipblas/hipblasLt 庫來實現。如何知道哪種實現最快並且應該被選擇？這就是 TunableOp 提供的功能。某些運算元已作為 Tunable Operators 使用多種策略實現。在執行時，會對所有策略進行效能分析，並選擇最快的一種用於所有後續操作。

請參閱 文件 以獲取如何使用它的資訊。

流檢查器 (原型)

CUDA 檢查器是一個用於檢測 PyTorch 中流之間同步錯誤的原型工具。請參閱 文件 以獲取如何使用它的資訊。

GPUDirect Storage (原型)

位於 torch.cuda.gds 中的 API 提供了某些 cuFile API 的輕量級包裝器，允許在 GPU 記憶體和儲存之間進行直接記憶體訪問傳輸，避免在 CPU 中使用跳板緩衝區 (bounce buffer)。請參閱 cufile API 文件 以獲取更多詳細資訊。

這些 API 可用於 CUDA 12.6 及更高版本。為了使用這些 API，必須確保系統已按照 GPUDirect Storage 文件 的說明進行了適當配置以使用 GPUDirect Storage。

請參閱 GdsFile 的文件以獲取如何使用這些 API 的示例。

gds_register_buffer	在 CUDA 裝置上將儲存註冊為 cufile 緩衝區。
gds_deregister_buffer	在 CUDA 裝置上登出先前已註冊的儲存作為 cufile 緩衝區。
GdsFile	cuFile 的包裝器。