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()

删除此实例当前持有的图。