Семантика подсчёта

• Подсчёты выполняются по формулам из flop_registry. Операторы без формулы могут быть разложены на зарегистрированные операторы; в противном случае они добавляют ноль FLOP.
• По умолчанию аргументы-тензоры в формулах заменяются их формами. Аргументы, не являющиеся тензорами, передаются без изменений. Используйте register_flop_formula(..., get_raw=True) только в том случае, если формуле нужны исходные аргументы-тензоры или метаданные.
• Операции прямого и обратного проходов учитываются, только если они выполняются во время активности менеджера контекста.
• В формулах по умолчанию используются плотные наивные математические определения. Например, умножение матриц учитывает 2 * m * n * k FLOP.
• Подсчёты автоматически не корректируются с учётом пропускной способности для конкретного типа данных, Tensor Core, разреженности, квантования, маскирования, пропущенных элементов, перемещения данных или зависящих от данных досрочных выходов. Формула должна явно учитывать такую семантику. Например, встроенные формулы attention задают верхние границы для каузального или иным образом маскированного attention, если формула явно не моделирует маску.
• Операторы высшего порядка и torch.compile могут по-разному представлять разложенные или объединённые вычисления по сравнению с eager-выполнением. Для пользовательских операторов и пользовательских ядер Triton нужна формула или разложение, если они должны учитываться в FLOP.
• Атрибуция модуля отслеживается с помощью ModuleTracker. Итоговые значения всегда доступны в "Global"; строки для подмодулей отражают атрибуцию вызовов модулей, наблюдавшихся во время работы контекста, по принципу best effort.
• Рабочая нагрузка по-прежнему выполняется в обычном режиме. Если важны накладные расходы или использование памяти, используйте этот режим для нескольких репрезентативных итераций, а не для длительных запусков обучения.

Пример

import torch
from torch.utils.flop_counter import FlopCounterMode

model = torch.nn.Linear(16, 32)
x = torch.randn(4, 16)

with FlopCounterMode(display=False) as mode:
    model(x).sum().backward()

print(mode.get_total_flops())
print(mode.get_flop_counts()["Global"])

Регистрация формулы для пользовательского оператора

Регистрируйте пользовательские формулы FLOP до создания FlopCounterMode. При инициализации режим сохраняет снимок глобального реестра.

from math import prod

import torch
from torch.utils.flop_counter import FlopCounterMode, register_flop_formula

@torch.library.custom_op("example::scale", mutates_args=())
def scale(x: torch.Tensor) -> torch.Tensor:
    return x * 2

@register_flop_formula(torch.ops.example.scale)
def scale_flops(x_shape, *, out_shape=None) -> int:
    return prod(x_shape)

x = torch.randn(8)
with FlopCounterMode(display=False) as mode:
    scale(x)

assert mode.get_total_flops() == 8