CUDAGraph

class torch.cuda.CUDAGraph

CUDA 图的包装器。

警告

此API目前处于测试阶段，可能会在未来的版本中进行更改。

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

开始在当前流上捕获CUDA工作。

通常情况下，你不应该自己调用 capture_begin。 使用 graph 或 make_graphed_callables()， 它们会在内部调用 capture_begin。

Parameters

• 池（可选）——令牌（由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)

Parameters

debug_path (必需) – 转储图的路径。

如果通过 CUDAGraph.enable_debug_mode() 启用了调试模式，则调用调试函数以转储图形。

enable_debug_mode()

启用CUDAGraph.debug_dump的调试模式。

pool()

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

这个id可以选择性地传递给另一个图的 capture_begin， 这提示另一个图可能共享相同的内存池。

replay()

重放由该图捕获的CUDA工作。

reset()

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