graph

class torch.cuda.graph(cuda_graph, pool=None, stream=None, capture_error_mode='global') [source]

Менеджер контекста, который записывает работу CUDA в объект torch.cuda.CUDAGraph для последующего воспроизведения.

См. Графы CUDA для общего введения, подробного использования и ограничений.

Параметры

• cuda_graph (torch.cuda.CUDAGraph) – Объект графа, используемый для захвата.
• pool (необязательно) – Непрозрачный токен (возвращаемый вызовом graph_pool_handle() или other_Graph_instance.pool()), указывающий, что захват этого графа может использовать память из указанного пула. См. Управление памятью графа.
• stream (torch.cuda.Stream, необязательно) – Если предоставлен, будет установлен в качестве текущего потока в контексте. Если не предоставлен, graph устанавливает собственный внутренний вспомогательный поток в качестве текущего потока в контексте.
• capture_error_mode (str, необязательно) – задает cudaStreamCaptureMode для потока захвата графа. Может быть «global», «thread_local» или «relaxed». Во время захвата графа CUDA некоторые действия, такие как cudaMalloc, могут быть небезопасными. «global» вызовет ошибку при действиях в других потоках, «thread_local» — только для действий в текущем потоке, а «relaxed» — без ошибки при действиях. НЕ изменяйте это значение, если вы не знакомы с cudaStreamCaptureMode

Примечание

Для эффективного совместного использования памяти, если вы передаёте pool , используемый предыдущим захватом, и предыдущий захват использовал явное аргумент stream, вы должны передать тот же аргумент stream этому захвату.

Предупреждение

Этот API находится в стадии бета-тестирования и может быть изменён в будущих выпусках.