FlopCounterMode

class torch.utils.flop_counter.FlopCounterMode(mods=None, depth=2, display=True, custom_mapping=None) [исходный код]

Подсчитывает теоретическое количество FLOP для операторов, выполняемых в контексте.

FlopCounterMode использует TorchDispatchMode для перехвата операций и применения зарегистрированных формул FLOP. Подсчёты основаны на формах тензоров и определяются формулами; неподдерживаемые операции дают нулевое количество FLOP, если только они не разлагаются на поддерживаемые операции.

Необязательный аргумент mods больше не требуется для атрибуции модулей. Если display имеет значение true, при выходе из контекста выводится таблица. Используйте get_total_flops() или get_flop_counts(), чтобы программно получить те же сведения.

Пример использования:

mod = ...
inp = ...

with FlopCounterMode(display=False) as flop_counter:
    mod(inp).sum().backward()

total = flop_counter.get_total_flops()

get_flop_counts() [исходный код]

Возвращает количество FLOP в виде словаря словарей.

Внешний словарь имеет ключи с именами модулей, а внутренний — ключи с именами операций.

Возвращает:

Количество FLOP в виде словаря.

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

Dict[str, Dict[Any, int]]