График

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

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

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

Параметры:

• cuda_graph (torch.cuda.CUDAGraph) – Объект графика, используемый для захвата.
• pool (необязательно) – Непрозрачный токен (возвращаемый вызовом graph_pool_handle() или other_Graph_instance.pool()), предполагающий, что захват этого графика может совместно использовать память из указанного пула. См. Управление памятью графика.
• stream (torch.cuda.Stream, необязательно) – Если указано, оно будет установлено как текущий поток в контексте. Если не указано, graph установит свой собственный внутренний побочный поток в качестве текущего потока в контексте.

Примечание

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

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

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