Shortcuts

Source code for mmengine.analysis.complexity_analysis

# Copyright (c) OpenMMLab. All rights reserved.

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(),
}


[docs]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__
[docs]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__
[docs]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()
[docs]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()
[docs]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
[docs]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()

© Copyright 2022, mmengine contributors. Revision 39ed23fa.

Built with Sphinx using a theme provided by Read the Docs.
Read the Docs v: latest
Versions
latest
stable
v0.10.3
v0.10.2
v0.10.1
v0.10.0
v0.9.1
v0.9.0
v0.8.5
v0.8.4
v0.8.3
v0.8.2
v0.8.1
v0.8.0
v0.7.4
v0.7.3
v0.7.2
v0.7.1
v0.7.0
v0.6.0
v0.5.0
v0.4.0
v0.3.0
v0.2.0
Downloads
epub
On Read the Docs
Project Home
Builds

Free document hosting provided by Read the Docs.