# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the license found in the
# LICENSE file in the root directory of this source tree.
import typing
from collections import defaultdict
from typing import Any, Counter, DefaultDict, Dict, Optional, Tuple, Union
import torch.nn as nn
from rich import box
from rich.console import Console
from rich.table import Table
from torch import Tensor
from .jit_analysis import JitModelAnalysis
from .jit_handles import (Handle, addmm_flop_jit, batchnorm_flop_jit,
bmm_flop_jit, conv_flop_jit, einsum_flop_jit,
elementwise_flop_counter, generic_activation_jit,
linear_flop_jit, matmul_flop_jit, norm_flop_counter)
# A dictionary that maps supported operations to their flop count jit handles.
_DEFAULT_SUPPORTED_FLOP_OPS: Dict[str, Handle] = {
'aten::addmm': addmm_flop_jit,
'aten::bmm': bmm_flop_jit,
'aten::_convolution': conv_flop_jit,
'aten::einsum': einsum_flop_jit,
'aten::matmul': matmul_flop_jit,
'aten::mm': matmul_flop_jit,
'aten::linear': linear_flop_jit,
# You might want to ignore BN flops due to inference-time fusion.
# Use `set_op_handle("aten::batch_norm", None)
'aten::batch_norm': batchnorm_flop_jit,
'aten::group_norm': norm_flop_counter(2),
'aten::layer_norm': norm_flop_counter(2),
'aten::instance_norm': norm_flop_counter(1),
'aten::upsample_nearest2d': elementwise_flop_counter(0, 1),
'aten::upsample_bilinear2d': elementwise_flop_counter(0, 4),
'aten::adaptive_avg_pool2d': elementwise_flop_counter(1, 0),
'aten::grid_sampler': elementwise_flop_counter(0, 4), # assume bilinear
}
# A dictionary that maps supported operations to
# their activation count handles.
_DEFAULT_SUPPORTED_ACT_OPS: Dict[str, Handle] = {
'aten::_convolution': generic_activation_jit('conv'),
'aten::addmm': generic_activation_jit(),
'aten::bmm': generic_activation_jit(),
'aten::einsum': generic_activation_jit(),
'aten::matmul': generic_activation_jit(),
'aten::linear': generic_activation_jit(),
}
class FlopAnalyzer(JitModelAnalysis):
"""Provides access to per-submodule model flop count obtained by tracing a
model with pytorch's jit tracing functionality.
By default, comes with standard flop counters for a few common operators.
Note:
- Flop is not a well-defined concept. We just produce our best
estimate.
- We count one fused multiply-add as one flop.
Handles for additional operators may be added, or the default ones
overwritten, using the ``.set_op_handle(name, func)`` method.
See the method documentation for details.
Flop counts can be obtained as:
- ``.total(module_name="")``: total flop count for the module
- ``.by_operator(module_name="")``: flop counts for the module, as a
Counter over different operator types
- ``.by_module()``: Counter of flop counts for all submodules
- ``.by_module_and_operator()``: dictionary indexed by descendant of
Counters over different operator types
An operator is treated as within a module if it is executed inside the
module's ``__call__`` method. Note that this does not include calls to
other methods of the module or explicit calls to ``module.forward(...)``.
Modified from
https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/flop_count.py
Args:
model (nn.Module): The model to analyze.
inputs (Union[Tensor, Tuple[Tensor, ...]]): The input to the model.
Examples:
>>> import torch.nn as nn
>>> import torch
>>> class TestModel(nn.Module):
... def __init__(self):
... super().__init__()
... self.fc = nn.Linear(in_features=1000, out_features=10)
... self.conv = nn.Conv2d(
... in_channels=3, out_channels=10, kernel_size=1
... )
... self.act = nn.ReLU()
... def forward(self, x):
... return self.fc(self.act(self.conv(x)).flatten(1))
>>> model = TestModel()
>>> inputs = (torch.randn((1,3,10,10)),)
>>> flops = FlopAnalyzer(model, inputs)
>>> flops.total()
13000
>>> flops.total("fc")
10000
>>> flops.by_operator()
Counter({"addmm" : 10000, "conv" : 3000})
>>> flops.by_module()
Counter({"" : 13000, "fc" : 10000, "conv" : 3000, "act" : 0})
>>> flops.by_module_and_operator()
{"" : Counter({"addmm" : 10000, "conv" : 3000}),
"fc" : Counter({"addmm" : 10000}),
"conv" : Counter({"conv" : 3000}),
"act" : Counter()
}
"""
def __init__(
self,
model: nn.Module,
inputs: Union[Tensor, Tuple[Tensor, ...]],
) -> None:
super().__init__(model=model, inputs=inputs)
self.set_op_handle(**_DEFAULT_SUPPORTED_FLOP_OPS)
__init__.__doc__ = JitModelAnalysis.__init__.__doc__
class ActivationAnalyzer(JitModelAnalysis):
"""Provides access to per-submodule model activation count obtained by
tracing a model with pytorch's jit tracing functionality.
By default, comes with standard activation counters for convolutional and
dot-product operators. Handles for additional operators may be added, or
the default ones overwritten, using the ``.set_op_handle(name, func)``
method. See the method documentation for details. Activation counts can be
obtained as:
- ``.total(module_name="")``: total activation count for a module
- ``.by_operator(module_name="")``: activation counts for the module,
as a Counter over different operator types
- ``.by_module()``: Counter of activation counts for all submodules
- ``.by_module_and_operator()``: dictionary indexed by descendant of
Counters over different operator types
An operator is treated as within a module if it is executed inside the
module's ``__call__`` method. Note that this does not include calls to
other methods of the module or explicit calls to ``module.forward(...)``.
Modified from
https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/activation_count.py
Args:
model (nn.Module): The model to analyze.
inputs (Union[Tensor, Tuple[Tensor, ...]]): The input to the model.
Examples:
>>> import torch.nn as nn
>>> import torch
>>> class TestModel(nn.Module):
... def __init__(self):
... super().__init__()
... self.fc = nn.Linear(in_features=1000, out_features=10)
... self.conv = nn.Conv2d(
... in_channels=3, out_channels=10, kernel_size=1
... )
... self.act = nn.ReLU()
... def forward(self, x):
... return self.fc(self.act(self.conv(x)).flatten(1))
>>> model = TestModel()
>>> inputs = (torch.randn((1,3,10,10)),)
>>> acts = ActivationAnalyzer(model, inputs)
>>> acts.total()
1010
>>> acts.total("fc")
10
>>> acts.by_operator()
Counter({"conv" : 1000, "addmm" : 10})
>>> acts.by_module()
Counter({"" : 1010, "fc" : 10, "conv" : 1000, "act" : 0})
>>> acts.by_module_and_operator()
{"" : Counter({"conv" : 1000, "addmm" : 10}),
"fc" : Counter({"addmm" : 10}),
"conv" : Counter({"conv" : 1000}),
"act" : Counter()
}
"""
def __init__(
self,
model: nn.Module,
inputs: Union[Tensor, Tuple[Tensor, ...]],
) -> None:
super().__init__(model=model, inputs=inputs)
self.set_op_handle(**_DEFAULT_SUPPORTED_ACT_OPS)
__init__.__doc__ = JitModelAnalysis.__init__.__doc__
def flop_count(
model: nn.Module,
inputs: Tuple[Any, ...],
supported_ops: Optional[Dict[str, Handle]] = None,
) -> Tuple[DefaultDict[str, float], Counter[str]]:
"""Given a model and an input to the model, compute the per-operator Gflops
of the given model.
Adopted from
https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/flop_count.py
Args:
model (nn.Module): The model to compute flop counts.
inputs (tuple): Inputs that are passed to `model` to count flops.
Inputs need to be in a tuple.
supported_ops (dict(str,Callable) or None) : provide additional
handlers for extra ops, or overwrite the existing handlers for
convolution and matmul and einsum. The key is operator name and
the value is a function that takes (inputs, outputs) of the op.
We count one Multiply-Add as one FLOP.
Returns:
tuple[defaultdict, Counter]: A dictionary that records the number of
gflops for each operation and a Counter that records the number of
unsupported operations.
"""
if supported_ops is None:
supported_ops = {}
flop_counter = FlopAnalyzer(model, inputs).set_op_handle(**supported_ops)
giga_flops = defaultdict(float)
for op, flop in flop_counter.by_operator().items():
giga_flops[op] = flop / 1e9
return giga_flops, flop_counter.unsupported_ops()
def activation_count(
model: nn.Module,
inputs: Tuple[Any, ...],
supported_ops: Optional[Dict[str, Handle]] = None,
) -> Tuple[DefaultDict[str, float], Counter[str]]:
"""Given a model and an input to the model, compute the total number of
activations of the model.
Adopted from
https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/activation_count.py
Args:
model (nn.Module): The model to compute activation counts.
inputs (tuple): Inputs that are passed to `model` to count activations.
Inputs need to be in a tuple.
supported_ops (dict(str,Callable) or None) : provide additional
handlers for extra ops, or overwrite the existing handlers for
convolution and matmul. The key is operator name and the value
is a function that takes (inputs, outputs) of the op.
Returns:
tuple[defaultdict, Counter]: A dictionary that records the number of
activation (mega) for each operation and a Counter that records the
number of unsupported operations.
"""
if supported_ops is None:
supported_ops = {}
act_counter = ActivationAnalyzer(model,
inputs).set_op_handle(**supported_ops)
mega_acts = defaultdict(float)
for op, act in act_counter.by_operator().items():
mega_acts[op] = act / 1e6
return mega_acts, act_counter.unsupported_ops()
def parameter_count(model: nn.Module) -> typing.DefaultDict[str, int]:
"""Count parameters of a model and its submodules.
Adopted from
https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/parameter_count.py
Args:
model (nn.Module): the model to count parameters.
Returns:
dict[str, int]: the key is either a parameter name or a module name.
The value is the number of elements in the parameter, or in all
parameters of the module. The key "" corresponds to the total
number of parameters of the model.
"""
count = defaultdict(int) # type: typing.DefaultDict[str, int]
for name, param in model.named_parameters():
size = param.numel()
name = name.split('.')
for k in range(0, len(name) + 1):
prefix = '.'.join(name[:k])
count[prefix] += size
return count
def parameter_count_table(model: nn.Module, max_depth: int = 3) -> str:
"""Format the parameter count of the model (and its submodules or
parameters)
Adopted from
https://github.com/facebookresearch/fvcore/blob/main/fvcore/nn/parameter_count.py
Args:
model (nn.Module): the model to count parameters.
max_depth (int): maximum depth to recursively print submodules or
parameters
Returns:
str: the table to be printed
"""
count: typing.DefaultDict[str, int] = parameter_count(model)
# pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
param_shape: typing.Dict[str, typing.Tuple] = {
k: tuple(v.shape)
for k, v in model.named_parameters()
}
# pyre-fixme[24]: Generic type `tuple` expects at least 1 type parameter.
rows: typing.List[typing.Tuple] = []
def format_size(x: int) -> str:
if x > 1e8:
return f'{x / 1e9:.1f}G'
if x > 1e5:
return f'{x / 1e6:.1f}M'
if x > 1e2:
return f'{x / 1e3:.1f}K'
return str(x)
def fill(lvl: int, prefix: str) -> None:
if lvl >= max_depth:
return
for name, v in count.items():
if name.count('.') == lvl and name.startswith(prefix):
indent = ' ' * (lvl + 1)
if name in param_shape:
rows.append(
(indent + name, indent + str(param_shape[name])))
else:
rows.append((indent + name, indent + format_size(v)))
fill(lvl + 1, name + '.')
rows.append(('model', format_size(count.pop(''))))
fill(0, '')
table = Table(
title=f'parameter count of {model.__class__.__name__}', box=box.ASCII2)
table.add_column('name')
table.add_column('#elements or shape')
for row in rows:
table.add_row(*row)
console = Console()
with console.capture() as capture:
console.print(table, end='')
return capture.get()