FlopCounterMode

class torch.utils.flop_counter.FlopCounterMode(mods=None, depth=2, display=True, custom_mapping=None)

Count theoretical FLOPs for operators that run inside the context.

FlopCounterMode uses TorchDispatchMode to intercept operations and apply registered FLOP formulas. Counts are shape-based and formula-defined; unsupported operations contribute zero FLOPs unless they decompose into supported operations.

The optional mods argument is no longer needed for module attribution. When display is true, exiting the context prints a table. Use get_total_flops() or get_flop_counts() to read the same information programmatically.

Example usage:

mod = ...
inp = ...

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

total = flop_counter.get_total_flops()

get_flop_counts()

Return the flop counts as a dictionary of dictionaries.

The outer dictionary is keyed by module name, and the inner dictionary is keyed by operation name.

Returns:

The flop counts as a dictionary.

Return type:

Dict[str, Dict[Any, int]]