CUDAGraph

class torch.cuda.CUDAGraph(keep_graph=False)

CUDA 圖的包裝器。

引數

keep_graph (bool, optional) – 如果 keep_graph=False，則 cudaGraphExec_t 將在 capture_end 結束時例項化在 GPU 上，並且底層的 cudaGraph_t 將被銷燬。希望在例項化之前查詢或以其他方式修改底層 cudaGraph_t 的使用者可以在 capture_end 之後設定 keep_graph=True 並透過 raw_cuda_graph 訪問它。請注意，在這種情況下，cudaGraphExec_t 不會在 capture_end 結束時例項化。相反，它將透過顯式呼叫 instantiate 來例項化，或者在首次呼叫 replay 時自動例項化（如果尚未呼叫 instantiate）。建議在手動呼叫 replay 之前手動呼叫 instantiate，以防止首次呼叫 replay 時延遲增加。允許在首次呼叫 instantiate 後修改原始 cudaGraph_t，但使用者必須再次手動呼叫 instantiate 以確保例項化的圖具有這些更改。PyTorch 無法跟蹤這些更改。

返回型別

自我

警告

此 API 處於 Beta 版，未來版本中可能會更改。

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

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

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

引數

• pool (optional) – Token（由 graph_pool_handle() 或 other_Graph_instance.pool() 返回）表示此圖可能與指示的池共享記憶體。請參閱 圖記憶體管理。

• capture_error_mode (str, optional) – 指定捕獲流的 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 (required) – 要將圖轉儲到的路徑。

如果透過 CUDAGraph.enable_debug_mode() 啟用了除錯，則呼叫除錯函式將圖轉儲到指定路徑。

enable_debug_mode()

啟用 CUDAGraph.debug_dump 的除錯模式。

instantiate()

例項化 CUDA 圖。如果 keep_graph=False，則由 capture_end 呼叫；如果 keep_graph=True 且尚未顯式呼叫 instantiate，則由 replay 呼叫。不會銷燬 raw_cuda_graph 返回的 cudaGraph_t。

pool()

返回一個不透明的 token，表示此圖的記憶體池的 ID。

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

返回型別

_POOL_HANDLE

raw_cuda_graph()

返回底層的 cudaGraph_t。 keep_graph 必須為 True。

有關如何操作此物件的 API，請參閱： Graph Managmement 和 cuda-python Graph Management 繫結

返回型別

int

raw_cuda_graph_exec()

返回底層的 cudaGraphExec_t。如果 keep_graph 為 True，則必須已呼叫 instantiate；如果 keep_graph 為 False，則必須已呼叫 capture_end。如果您在呼叫 raw_cuda_graph_exec() 後呼叫 instantiate()，則之前返回的 cudaGraphExec_t 將被銷燬。您有責任不要在此物件被銷燬後使用它。

有關如何操作此物件的 API，請參閱： Graph Execution 和 cuda-python Graph Execution 繫結

返回型別

int

replay()

重放此圖捕獲的 CUDA 工作。

reset()

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