graph

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

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

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

Параметры:

• cuda_graph (torch.cuda.CUDAGraph) – Объект графа, используемый для захвата.
• pool (необязательный) – Непрозрачный токен (возвращаемый вызовом graph_pool_handle() или other_Graph_instance.pool()) либо MemPool, указывающий, что при захвате этого графа может использоваться общая память из указанного пула. См. раздел Управление памятью графов.
• stream (torch.cuda.Stream, необязательный) – Если задан, он будет установлен как текущий поток в контексте. Если не задан, graph устанавливает собственный внутренний вспомогательный поток как текущий поток в контексте.
• capture_error_mode (str, необязательный) – задаёт cudaStreamCaptureMode для потока захвата графа. Возможные значения: “global”, “thread_local” или “relaxed”. Во время захвата графа CUDA некоторые действия, например cudaMalloc, могут быть небезопасными. При значении “global” ошибка будет возникать при выполнении таких действий в других потоках, при “thread_local” — только при выполнении таких действий в текущем потоке, а при “relaxed” ошибки при выполнении таких действий возникать не будут. НЕ меняйте эту настройку, если вы не знакомы с cudaStreamCaptureMode
• enable_annotations (bool, необязательный) – Если True, включает запись аннотаций ядер при входе и автоматически вызывает resolve_pending_annotations() перед завершением захвата. При выходе аннотации не очищаются, поэтому в одной рабочей нагрузке можно накапливать аннотации для нескольких графов. Требуется пакет cuda.bindings и cuda-compat версии 13.1 или выше либо драйвер CUDA версии 13.1 или выше.

• check_input_liveness (bool, необязательный) –

  Если True, отслеживает внешние входные тензоры во время захвата графа и вызывает ошибку, если какой-либо из них освобождается до воспроизведения. Это помогает выявлять ошибки «использование после освобождения», возникающие, когда сборщик мусора удаляет входные тензоры между захватом и воспроизведением. По умолчанию: False.

  Примечание

  Пользовательские ядра CUDA, добавленные вне PyTorch (например, с помощью cuLaunchKernel или DLPack), этим механизмом не отслеживаются.

Примечание

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

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

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