CUDAGraph

class torch.cuda.CUDAGraph

CUDA 圖的包裝器。

警告

此 API 處於 Beta 階段，將來版本中可能會發生變化。

capture_begin(pool=None, capture_error_mode='global')

開始在當前流上捕獲 CUDA 工作。

通常情況下，您不應自行呼叫 capture_begin。請使用 graph 或 make_graphed_callables()，它們在內部呼叫 capture_begin。

引數

• pool (可選) – 令牌（由 torch.cuda.graph_pool_handle() 或 other_Graph_instance.pool() 返回），提示此圖可能與指定的池共享記憶體。參見圖記憶體管理。

• capture_error_mode (str, 可選) – 指定圖捕獲流的 cudaStreamCaptureMode。可以是 “global”、“thread_local” 或 “relaxed”。在 CUDA 圖捕獲期間，某些操作（例如 cudaMalloc）可能不安全。“global” 會對其他執行緒中的操作報錯，“thread_local” 只會針對當前執行緒中的操作報錯，“relaxed” 不會針對這些操作報錯。除非您熟悉 cudaStreamCaptureMode，否則請勿更改此設定。

capture_end()

結束在當前流上的 CUDA 圖捕獲。

在 capture_end 之後，可以在此例項上呼叫 replay。

通常情況下，您不應自行呼叫 capture_end。請使用 graph 或 make_graphed_callables()，它們在內部呼叫 capture_end。

debug_dump(debug_path)

引數

debug_path (必填) – 圖的轉儲路徑。

如果透過 CUDAGraph.enable_debug_mode() 啟用了除錯模式，則呼叫除錯函式轉儲圖。

enable_debug_mode()

為 CUDAGraph.debug_dump 啟用除錯模式。

pool()

返回一個不透明令牌，代表此圖記憶體池的 ID。

此 ID 可選擇傳遞給另一個圖的 capture_begin，提示另一個圖可能共享相同的記憶體池。

replay()

重放此圖捕獲的 CUDA 工作。

reset()

刪除此例項當前持有的圖。