CUDAGraph

class torch.cuda.CUDAGraph[source][source]

CUDA 图包装器。

警告

此 API 处于测试阶段，未来版本中可能会有所变化。

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

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

通常，您不应该直接调用 capture_begin 。请使用 graph 或 make_graphed_callables() ，它们会内部调用 capture_begin 。

参数:

• pool（可选）- 由 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()[source][source]

在当前流上结束 CUDA 图捕获。

在 capture_end 之后， replay 可能被调用在此实例上。

通常情况下，您不应该直接调用 capture_end 。请使用 graph 或 make_graphed_callables() ，它们会内部调用 capture_end 。

debug_dump(debug_path)[source][source]

参数:

debug_path（必需）- 存储图的路径。

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

启用调试模式()[source][source] ¶

为 CUDAGraph.debug_dump 启用调试模式。

pool()

返回表示此图内存池 id 的不可见令牌。

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

replay()

回放此图捕获的 CUDA 工作。

reset()

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