CUDAGraph

类torch.cuda.CUDAGraph[源代码]

CUDA图的封装类。

警告

此 API 处于 beta 阶段，未来版本可能有所更改。

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()[源代码]

结束当前流的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 的调试模式。

pool()[源代码]

返回一个代表该图形内存池ID的不透明令牌。

此 ID 可以选择性地传递给另一个图的 capture_begin 方法，暗示该图可能会使用同一个内存池。

replay()[源代码]

重放该图捕捉到的CUDA任务。

reset()[源代码]

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