torch.cuda

這個套件新增了對 CUDA 張量類型的支援。

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

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

CUDA 語意 提供了更多關於使用 CUDA 的細節。

StreamContext	選擇給定 stream 的 context-manager。
can_device_access_peer	檢查兩個裝置之間是否可以進行 peer access。
current_blas_handle	回傳指向目前 cuBLAS handle 的 cublasHandle_t 指標
current_device	回傳目前所選裝置的索引。
current_stream	回傳給定裝置目前所選的 Stream。
cudart	檢索 CUDA runtime API 模組。
default_stream	回傳給定裝置的預設 Stream。
device	變更所選裝置的 context-manager。
device_count	回傳可用的 GPU 數量。
device_memory_used	以位元組為單位回傳 nvidia-smi 或 amd-smi 所提供的已使用全域（裝置）記憶體。
device_of	將目前裝置變更為給定物件的裝置的 context-manager。
get_arch_list	回傳此程式庫編譯時所針對的 CUDA 架構列表。
get_device_capability	取得裝置的 CUDA 功能。
get_device_name	取得裝置的名稱。
get_device_properties	取得裝置的屬性。
get_gencode_flags	回傳此程式庫編譯時所使用的 NVCC gencode 標誌。
get_sync_debug_mode	回傳 CUDA 同步操作的除錯模式的目前值。
init	初始化 PyTorch 的 CUDA 狀態。
ipc_collect	強制回收 CUDA IPC 釋放後的 GPU 記憶體。
is_available	回傳一個布林值，指示目前 CUDA 是否可用。
is_initialized	回傳 PyTorch 的 CUDA 狀態是否已初始化。
memory_usage	回傳在過去的取樣期間內，全域（裝置）記憶體正在讀取或寫入的時間百分比，由 nvidia-smi 提供。
set_device	設定目前的裝置。
set_stream	設定目前的 stream。這是一個用於設定 stream 的包裝 API。
set_sync_debug_mode	設定 CUDA 同步操作的除錯模式。
stream	包裝選擇給定 stream 的 Context-manager StreamContext。
synchronize	等待 CUDA 裝置上所有 stream 中的所有核心完成。
utilization	回傳在過去的取樣期間內，一個或多個核心正在 GPU 上執行的時間百分比，由 nvidia-smi 提供。
temperature	回傳 GPU 感測器的平均溫度，單位為攝氏度（Centigrades）。
power_draw	回傳 GPU 感測器的平均功耗，單位為 mW (MilliWatts)。
clock_rate	回傳在過去的取樣期間內，GPU SM 的時脈速度，單位為赫茲 (Hz)，由 nvidia-smi 提供。
OutOfMemoryError	當裝置記憶體不足時引發的例外。

隨機數產生器

get_rng_state	以 ByteTensor 的形式回傳指定 GPU 的隨機數產生器狀態。
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 裝置收集張量。

Streams 和 Events

Stream	CUDA stream 的包裝。
ExternalStream	外部配置的 CUDA stream 的包裝。
Event	CUDA event 的包裝。

Graphs (beta)

is_current_stream_capturing	如果目前 CUDA stream 上正在進行 CUDA graph 擷取，則回傳 True，否則回傳 False。
graph_pool_handle	回傳表示 graph 記憶體池 ID 的不透明權杖。
CUDAGraph	CUDA graph 的包裝。
graph	Context-manager，將 CUDA 工作擷取到 torch.cuda.CUDAGraph 物件中，以便稍後重播。
make_graphed_callables	接受可呼叫物件（函數或 nn.Modules）並回傳 graph 版本。

記憶體管理

empty_cache	釋放快取分配器目前持有的所有未佔用快取記憶體，以便這些記憶體可以在其他 GPU 應用程式中使用，並在 nvidia-smi 中可見。
get_per_process_memory_fraction	取得程序的記憶體比例。
list_gpu_processes	回傳給定裝置的執行程序及其 GPU 記憶體使用情況的人類可讀列印輸出。
mem_get_info	使用 cudaMemGetInfo 回傳給定裝置的全域可用和總 GPU 記憶體。
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 記憶體配置器追蹤的「峰值」統計資訊。
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 or int, optional) – 選定的裝置。如果 device 是 None (預設)，則使用 current_device() 給定的目前裝置上的 MemPool。

NVIDIA 工具擴展 (NVTX)

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

Jiterator（測試版）

jiterator._create_jit_fn	為 elementwise 運算建立 jiterator 產生的 cuda 核心。
jiterator._create_multi_output_jit_fn	為 elementwise 運算建立 jiterator 產生的 cuda 核心，該核心支援傳回一個或多個輸出。

TunableOp

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

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

Stream Sanitizer (原型)

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