快捷方式

CUDAGraph

class torch.cuda.CUDAGraph[源代码]

CUDA 图的包装器。

警告

此 API 处于测试阶段,可能会在将来的版本中发生更改。

capture_begin(pool=None, capture_error_mode='global')[源代码]

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

通常,您不应该自己调用 capture_begin。使用 graphmake_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。使用 graphmake_graphed_callables(),它们会在内部调用 capture_end

debug_dump(debug_path)[源代码]
参数

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

如果通过 CUDAGraph.enable_debug_mode() 启用了调试,则调用调试函数来转储图

enable_debug_mode()[源代码]

为 CUDAGraph.debug_dump 启用调试模式。

pool()[源代码]

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

此 ID 可以选择性地传递到另一个图的 capture_begin,这提示另一个图可能共享相同的内存池。

replay()[源代码]

重放此图捕获的 CUDA 工作。

reset()[source]

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

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取针对初学者和高级开发人员的深度教程

查看教程

资源

查找开发资源并获取问题的解答

查看资源