Add experimental aten backend

sovrasov · sovrasov · commit fec84af4a639 · 2024-01-21T15:12:34.000+09:00
diff --git a/ptflops/__init__.py b/ptflops/__init__.py
@@ -1,17 +1,18 @@
 '''
-Copyright (C) 2019-2023 Sovrasov V. - All Rights Reserved
+Copyright (C) 2019-2024 Sovrasov V. - All Rights Reserved
  * You may use, distribute and modify this code under the
  * terms of the MIT license.
  * You should have received a copy of the MIT license with
  * this file. If not visit https://opensource.org/licenses/MIT
 '''
 
 
-from .flops_counter import get_model_complexity_info
+from .flops_counter import FLOPS_BACKEND, get_model_complexity_info
 from .utils import flops_to_string, params_to_string
 
 __all__ = [
     "get_model_complexity_info",
     "flops_to_string",
     "params_to_string",
+    "FLOPS_BACKEND",
     ]
diff --git a/ptflops/aten_engine.py b/ptflops/aten_engine.py
@@ -0,0 +1,112 @@
+'''
+Copyright (C) 2024 Sovrasov V. - All Rights Reserved
+ * You may use, distribute and modify this code under the
+ * terms of the MIT license.
+ * You should have received a copy of the MIT license with
+ * this file. If not visit https://opensource.org/licenses/MIT
+'''
+
+
+import sys
+import traceback
+from collections import defaultdict
+from typing import Optional, Tuple, Union
+
+import torch
+from torch.utils._python_dispatch import TorchDispatchMode
+
+from ptflops.pytorch_engine import get_model_parameters_number
+from .aten_ops import ATEN_OPS_MAPPING
+
+
+def normalize_tuple(x):
+    if not isinstance(x, tuple):
+        return (x,)
+    return x
+
+
+class FlopCounterMode(TorchDispatchMode):
+    def __init__(self, module=None):
+        self.flop_counts = defaultdict(lambda: defaultdict(int))
+        self.parents = ['Global']
+        self._total_complexity = None
+        if module is not None:
+            for name, mod in dict(module.named_children()).items():
+                mod.register_forward_pre_hook(self.enter_module(name))
+                mod.register_forward_hook(self.exit_module(name))
+
+    @property
+    def complexity(self):
+        return self._total_complexity
+
+    def enter_module(self, name):
+        def f(*args):
+            self.parents.append(name)
+        return f
+
+    def exit_module(self, name):
+        def f(*args):
+            assert(self.parents[-1] == name)
+            self.parents.pop()
+        return f
+
+    def __enter__(self):
+        self.flop_counts.clear()
+        super().__enter__()
+
+    def __exit__(self, *args):
+        self._total_complexity = sum(self.flop_counts['Global'].values())
+        super().__exit__(*args)
+
+    def __torch_dispatch__(self, func, types, args=(), kwargs=None):
+        kwargs = kwargs if kwargs else {}
+
+        out = func(*args, **kwargs)
+        func_packet = func._overloadpacket
+        if func_packet in ATEN_OPS_MAPPING:
+            flop_count = ATEN_OPS_MAPPING[func_packet](args, normalize_tuple(out))
+            for par in self.parents:
+                self.flop_counts[par][func_packet] += flop_count
+
+        return out
+
+
+def get_flops_aten(model, input_res,
+                   print_per_layer_stat=True,
+                   input_constructor=None, ost=sys.stdout,
+                   verbose=False, ignore_modules=[],
+                   custom_modules_hooks={},
+                   output_precision=2,
+                   flops_units: Optional[str] = 'GMac',
+                   param_units: Optional[str] = 'M') -> Tuple[Union[int, None],
+                                                              Union[int, None]]:
+
+    params_sum = get_model_parameters_number(model)
+
+    if input_constructor:
+        batch = input_constructor(input_res)
+    else:
+        try:
+            batch = torch.ones(()).new_empty((1, *input_res),
+                                             dtype=next(model.parameters()).dtype,
+                                             device=next(model.parameters()).device)
+        except StopIteration:
+            batch = torch.ones(()).new_empty((1, *input_res))
+
+    try:
+        counter = FlopCounterMode(model)
+        with counter:
+            if isinstance(batch, dict):
+                _ = model(**batch)
+            else:
+                _ = model(batch)
+        macs_count = counter.complexity
+
+    except Exception as e:
+        print("Flops estimation was not finished successfully because of"
+              f"the following exception:\n{type(e)} : {e}")
+        traceback.print_exc()
+
+        return None, None
+
+    return macs_count, params_sum
diff --git a/ptflops/aten_ops.py b/ptflops/aten_ops.py
@@ -0,0 +1,118 @@
+'''
+Copyright (C) 2023 Sovrasov V. - All Rights Reserved
+ * You may use, distribute and modify this code under the
+ * terms of the MIT license.
+ * You should have received a copy of the MIT license with
+ * this file. If not visit https://opensource.org/licenses/MIT
+'''
+
+from typing import Any, List
+
+import torch
+
+aten = torch.ops.aten
+
+
+def get_shape(i):
+    return i.shape
+
+
+def prod(x):
+    res = 1
+    for i in x:
+        res *= i
+    return res
+
+
+def matmul_flop(inputs: List[Any], outputs: List[Any]) -> int:
+    """
+    Count flops for matmul.
+    """
+    # Inputs should be a list of length 2.
+    # Inputs contains the shapes of two matrices.
+    input_shapes = [get_shape(v) for v in inputs]
+    assert len(input_shapes) == 2, input_shapes
+    assert input_shapes[0][-1] == input_shapes[1][-2], input_shapes
+    flop = prod(input_shapes[0]) * input_shapes[-1][-1]
+    return flop
+
+
+def addmm_flop(inputs: List[Any], outputs: List[Any]) -> int:
+    """
+    Count flops for fully connected layers.
+    """
+    # Count flop for nn.Linear
+    # inputs is a list of length 3.
+    input_shapes = [get_shape(v) for v in inputs[1:3]]
+    # input_shapes[0]: [batch size, input feature dimension]
+    # input_shapes[1]: [batch size, output feature dimension]
+    assert len(input_shapes[0]) == 2, input_shapes[0]
+    assert len(input_shapes[1]) == 2, input_shapes[1]
+    batch_size, input_dim = input_shapes[0]
+    output_dim = input_shapes[1][1]
+    flops = batch_size * input_dim * output_dim
+    return flops
+
+
+def bmm_flop(inputs: List[Any], outputs: List[Any]) -> int:
+    """
+    Count flops for the bmm operation.
+    """
+    # Inputs should be a list of length 2.
+    # Inputs contains the shapes of two tensor.
+    assert len(inputs) == 2, len(inputs)
+    input_shapes = [get_shape(v) for v in inputs]
+    n, c, t = input_shapes[0]
+    d = input_shapes[-1][-1]
+    flop = n * c * t * d
+    return flop
+
+
+def conv_flop_count(
+    x_shape: List[int],
+    w_shape: List[int],
+    out_shape: List[int],
+    transposed: bool = False,
+) -> int:
+    """
+    Count flops for convolution. Note only multiplication is
+    counted. Computation for addition and bias is ignored.
+    Flops for a transposed convolution are calculated as
+    flops = (x_shape[2:] * prod(w_shape) * batch_size).
+    Args:
+        x_shape (list(int)): The input shape before convolution.
+        w_shape (list(int)): The filter shape.
+        out_shape (list(int)): The output shape after convolution.
+        transposed (bool): is the convolution transposed
+    Returns:
+        int: the number of flops
+    """
+    batch_size = x_shape[0]
+    conv_shape = (x_shape if transposed else out_shape)[2:]
+    flop = batch_size * prod(w_shape) * prod(conv_shape)
+    return flop
+
+
+def conv_flop(inputs: List[Any], outputs: List[Any]):
+    """
+    Count flops for convolution.
+    """
+    x, w = inputs[:2]
+    x_shape, w_shape, out_shape = (get_shape(x), get_shape(w), get_shape(outputs[0]))
+    transposed = inputs[6]
+
+    return conv_flop_count(x_shape, w_shape, out_shape, transposed=transposed)
+
+
+def transpose_shape(shape):
+    return [shape[1], shape[0]] + list(shape[2:])
+
+
+ATEN_OPS_MAPPING = {
+    aten.mm: matmul_flop,
+    aten.matmul: matmul_flop,
+    aten.addmm: addmm_flop,
+    aten.bmm: bmm_flop,
+    aten.convolution: conv_flop,
+    aten._convolution: conv_flop,
+}
diff --git a/ptflops/flops_counter.py b/ptflops/flops_counter.py
@@ -1,20 +1,27 @@
 '''
-Copyright (C) 2019-2023 Sovrasov V. - All Rights Reserved
+Copyright (C) 2019-2024 Sovrasov V. - All Rights Reserved
  * You may use, distribute and modify this code under the
  * terms of the MIT license.
  * You should have received a copy of the MIT license with
  * this file. If not visit https://opensource.org/licenses/MIT
 '''
 
 import sys
+from enum import Enum
 from typing import Any, Callable, Dict, List, Optional, TextIO, Tuple, Union
 
 import torch.nn as nn
 
+from .aten_engine import get_flops_aten
 from .pytorch_engine import get_flops_pytorch
 from .utils import flops_to_string, params_to_string
 
 
+class FLOPS_BACKEND(Enum):
+    PYTORCH = 'pytorch'
+    ATEN = 'aten'
+
+
 def get_model_complexity_info(model: nn.Module,
                               input_res: Tuple[int, ...],
                               print_per_layer_stat: bool = True,
@@ -24,7 +31,7 @@ def get_model_complexity_info(model: nn.Module,
                               verbose: bool = False,
                               ignore_modules: List[nn.Module] = [],
                               custom_modules_hooks: Dict[nn.Module, Any] = {},
-                              backend: str = 'pytorch',
+                              backend: Union[str, FLOPS_BACKEND] = FLOPS_BACKEND.PYTORCH,
                               flops_units: Optional[str] = None,
                               param_units: Optional[str] = None,
                               output_precision: int = 2) -> Tuple[Union[str, int, None],
@@ -58,6 +65,8 @@ def get_model_complexity_info(model: nn.Module,
     :type ignore_modules: nn.Module
     :param custom_modules_hooks: A dict that contains custom hooks on torch modules.
     :type custom_modules_hooks: Dict[nn.Module, Any]
+    :param backend: Backend that used for evaluating model complexity.
+    :type backend: FLOPS_BACKEND
     :param flops_units: Units for string representation of MACs (GMac, MMac or KMac).
     :type flops_units: Optional[str]
     :param param_units: Units for string representation of params (M, K or B).
@@ -75,7 +84,7 @@ def get_model_complexity_info(model: nn.Module,
     assert len(input_res) >= 1
     assert isinstance(model, nn.Module)
 
-    if backend == 'pytorch':
+    if FLOPS_BACKEND(backend) == FLOPS_BACKEND.PYTORCH:
         flops_count, params_count = get_flops_pytorch(model, input_res,
                                                       print_per_layer_stat,
                                                       input_constructor, ost,
@@ -84,6 +93,15 @@ def get_model_complexity_info(model: nn.Module,
                                                       output_precision=output_precision,
                                                       flops_units=flops_units,
                                                       param_units=param_units)
+    elif FLOPS_BACKEND(backend) == FLOPS_BACKEND.ATEN:
+        flops_count, params_count = get_flops_aten(model, input_res,
+                                                   print_per_layer_stat,
+                                                   input_constructor, ost,
+                                                   verbose, ignore_modules,
+                                                   custom_modules_hooks,
+                                                   output_precision=output_precision,
+                                                   flops_units=flops_units,
+                                                   param_units=param_units)
     else:
         raise ValueError('Wrong backend name')
 
