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，可以将 keep_graph=True 设置为 True，并在 capture_end 之后通过 raw_cuda_graph 进行访问。注意，在这种情况下，cudaGraphExec_t 不会在 capture_end 结束时实例化。相反，它将通过显式调用 instantiate 进行实例化，或者如果之前未调用 instantiate，则在第一次调用 replay 时自动实例化。建议在 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) – 令牌（由 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()

返回一个不透明令牌，表示此图的内存池 ID。

此 ID 可以选择性地传递给另一个图的 capture_begin，这暗示着其他图可能共享相同的内存池。

返回类型:

_POOL_HANDLE

raw_cuda_graph()

返回底层的 cudaGraph_t。keep_graph 必须为 True。

有关如何操作此对象的 API，请参阅：图管理 (Graph Management) 和 cuda-python 图管理绑定。

返回类型:

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 图执行绑定。

返回类型:

int

replay()

重放该图捕获的 CUDA 工作。

reset()

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