CUDAGraph¶
- class torch.cuda.CUDAGraph[源代码]¶
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()[源代码]¶
结束当前流上的 CUDA 图捕获。
在
capture_end
之后,可以对该实例调用replay
。通常,您不应该自己调用
capture_end
。使用graph
或make_graphed_callables()
,它们会在内部调用capture_end
。