torch.cuda

創建於: 2016年12月23日 | 最後更新於: 2025年08月19日

此包增加了對 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 裝置是否支援 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 感測器平均溫度（攝氏度）。
power_draw	返回 GPU 感測器平均功耗（毫瓦）。
clock_rate	返回過去取樣週期內 GPU SM 的時鐘速度（MHz），如“nvidia-smi”所示。
AcceleratorError	在裝置上執行時引發的異常
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.reduce_add_coalesced	從多個 GPU 求和張量。
comm.scatter	將張量散佈到多個 GPU。
comm.gather	從多個 GPU 裝置收集張量。

流和事件

流	CUDA 流的包裝器。
ExternalStream	外部分配的 CUDA 流的包裝器。
事件	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）並返回圖化後的版本。

此包增加了對 CUDA 中實現的裝置記憶體管理的支援。

記憶體管理

empty_cache	釋放快取分配器當前持有的所有未佔用快取記憶體，以便它們可供其他 GPU 應用程式使用並在“nvidia-smi”中可見。
get_per_process_memory_fraction	獲取程序的記憶體比例。
list_gpu_processes	返回給定裝置的正在執行的程序及其 GPU 記憶體使用情況的可讀列印輸出。
mem_get_info	使用 cudaMemGetInfo 返回給定裝置的全域性空閒和總 GPU 記憶體。
memory_stats	返回給定裝置的 CUDA 記憶體分配器統計資訊的字典。
memory_stats_as_nested_dict	以巢狀字典的形式返回 memory_stats() 的結果。
reset_accumulated_memory_stats	重置 CUDA 記憶體分配器跟蹤的“累積”（歷史）統計資訊。
host_memory_stats	返回給定裝置的 CUDA 記憶體分配器統計資訊的字典。
host_memory_stats_as_nested_dict	將“host_memory_stats()”的結果作為巢狀字典返回。
reset_accumulated_host_memory_stats	重置主機記憶體分配器跟蹤的“累積”（歷史）統計資訊。
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 代表快取分配器中的記憶體池。

caching_allocator_enable

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

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

一個上下文管理器，可將分配路由到給定的池。

引數

• pool (torch.cuda.MemPool) – 一個 MemPool 物件，將使其成為活動狀態，以便分配路由到此池。

• device (torch.device 或 int, optional) – 選定的裝置。如果 device 為 None（預設），則使用當前裝置上的 MemPool，由 current_device() 提供。

注意

此上下文管理器僅使當前執行緒的分配路由到給定的池。如果在上下文管理器內生成新執行緒（例如，透過呼叫 backward），則該執行緒中的分配不會路由到給定的池。

NVIDIA Tools Extension (NVTX)

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

Jiterator (beta)

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

TunableOp

某些操作可以使用多個庫或多種技術來實現。例如，GEMM 可以使用 cublas/cublasLt 庫或 hipblas/hipblasLt 庫分別針對 CUDA 或 ROCm 實現。如何知道哪種實現最快並應被選擇？這正是 TunableOp 提供的。某些運算子已透過 Tunable Operators 的形式實現。在執行時，所有策略都將進行效能分析，然後選擇最快的一種用於所有後續操作。

有關如何使用它的資訊，請參閱文件。

Stream Sanitizer (prototype)

CUDA Sanitizer 是一個用於檢測 PyTorch 中流之間同步錯誤的prototype工具。有關如何使用它的資訊，請參閱文件。

GPUDirect Storage (prototype)

torch.cuda.gds 中的 API 是對某些 cuFile API 的薄封裝，它們允許 GPU 記憶體和儲存之間的直接記憶體訪問傳輸，避免了 CPU 中的反彈緩衝區。有關更多詳細資訊，請參閱cufile API 文件。

這些 API 可在 CUDA 12.6 及更高版本中使用。為了使用這些 API，必須確保您的系統已根據GPUDirect Storage 文件正確配置為使用 GPUDirect Storage。

有關如何使用這些 API 的示例，請參閱 GdsFile 的文件。

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