torch.cuda.graph_annotations.get_kernel_annotations

torch.cuda.graph_annotations.get_kernel_annotations() → Mapping[int, list] [исходный код]

Возвращает текущий реестр записанных аннотаций ядер.

Ключи — непрозрачные целые числа, соответствующие полю graph node id, которое профилировщики на основе CUPTI добавляют к событиям ядер; значения — списки словарей аннотаций, записанных для этого узла. Реестр накапливает данные между захватами и является глобальным для процесса.

Возвращаемое отображение — это динамическое представление: оно обновляется на месте при инстанцировании графа (ключи аннотаций переназначаются на идентификаторы исполняемого графа), поэтому ссылка, полученная заранее, остаётся актуальной. Ключи можно использовать для сопоставления с трассировкой профилировщика после инстанцирования соответствующих графов. Рассматривайте отображение как доступное только для чтения; создайте его снимок с помощью dict(...), если нужна изоляция.

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

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

Пример:

>>> annotations = torch.cuda.graph_annotations.get_kernel_annotations()
>>> with open("annotations.pkl", "wb") as f:
...     pickle.dump(dict(annotations), f)

Тип возвращаемого значения:

Mapping[int, list[Any]]