Add an option to ignore funtionals in torch backend

Author: sovrasov

CHANGELOG.md

@@ -1,5 +1,10 @@
 # ptflops versions log
 
+## v 0.7.4
+- Switch to aten by default.
+- Add ignore and custom modules for aten.
+- Add an option to disable counting of functional-style operations in pytorch backend.
+
 ## v 0.7.3
 - Add aten backend to collect the amount of flops on aten level.
 

README.md

@@ -9,6 +9,7 @@ print per-layer computational cost of a given network.
 `ptflops` has two backends, `pytorch` and `aten`. `pytorch` backend is a legacy one, it considers `nn.Modules` only. However,
 it's still useful, since it provides a better par-layer analytics for CNNs. In all other cases it's recommended to use
 `aten` backend, which considers aten operations, and therefore it covers more model architectures (including transformers).
+The default backend is `aten`. Please, don't use `pytorch` backend for transformer architectures.
 
 ## `aten` backend
 ### Operations considered:
@@ -19,6 +20,9 @@ it's still useful, since it provides a better par-layer analytics for CNNs. In a
 - Use `verbose=True` to see the operations which were not considered during complexity computation.
 - This backend prints per-module statistics only for modules directly nested into the root `nn.Module`.
 Deeper modules at the second level of nesting are not shown in the per-layer statistics.
+- `ignore_modules` option forces `ptflops` to ignore the listed modules. This can be useful
+for research purposes. For instance, one can drop all convolutions from the counting process
+specifying `ignore_modules=[torch.ops.aten.convolution, torch.ops.aten._convolution]`.
 
 ## `pytorch` backend
 ### Supported layers:
@@ -41,7 +45,9 @@ Experimental support:
 
 - This backend doesn't take into account some of the `torch.nn.functional.*` and `tensor.*` operations. Therefore unsupported operations are
 not contributing to the final complexity estimation. See `ptflops/pytorch_ops.py:FUNCTIONAL_MAPPING,TENSOR_OPS_MAPPING` to check supported ops.
-- `ptflops` launches a given model on a random tensor and estimates amount of computations during inference. Complicated models can have several inputs, some of them could be optional. To construct non-trivial input one can use the `input_constructor` argument of the `get_model_complexity_info`. `input_constructor` is a function that takes the input spatial resolution as a tuple and returns a dict with named input arguments of the model. Next this dict would be passed to the model as a keyword arguments.
+Sometimes considering functional style conflicts with hooks for `nn.Module` (for instance, custom ones). In that case, counting with these ops can be disabled by
+passing `backend_specific_config={"count_functional" : False}`.
+- `ptflops` launches a given model on a random tensor and estimates amount of computations during inference. Complicated models can have several inputs, some of them could be optional. To construct non-trivial input one can use the `input_constructor` argument of the `get_model_complexity_info`. `input_constructor` is a function that takes the input spatial resolution as a tuple and returns a dict with named input arguments of the model. Next, this dict would be passed to the model as a keyword arguments.
 - `verbose` parameter allows to get information about modules that don't contribute to the final numbers.
 - `ignore_modules` option forces `ptflops` to ignore the listed modules. This can be useful
 for research purposes. For instance, one can drop all convolutions from the counting process

ptflops/aten_engine.py

@@ -12,7 +12,7 @@
 from collections import defaultdict
 from copy import deepcopy
 from functools import partial
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 from torch.utils._python_dispatch import TorchDispatchMode
@@ -106,8 +106,9 @@ def get_flops_aten(model, input_res,
                    custom_modules_hooks={},
                    output_precision=2,
                    flops_units: Optional[str] = 'GMac',
-                   param_units: Optional[str] = 'M') -> Tuple[Union[int, None],
-                                                              Union[int, None]]:
+                   param_units: Optional[str] = 'M',
+                   extra_config: Dict = {}) -> Tuple[Union[int, None],
+                                                     Union[int, None]]:
 
     params_sum = get_model_parameters_number(model)
     model.eval()

ptflops/flops_counter.py

@@ -34,8 +34,10 @@ def get_model_complexity_info(model: nn.Module,
                               backend: Union[str, FLOPS_BACKEND] = FLOPS_BACKEND.ATEN,
                               flops_units: Optional[str] = None,
                               param_units: Optional[str] = None,
-                              output_precision: int = 2) -> Tuple[Union[str, int, None],
-                                                                  Union[str, int, None]]:
+                              output_precision: int = 2,
+                              backend_specific_config: Dict = {}) -> Tuple[
+                                  Union[str, int, None],
+                                  Union[str, int, None]]:
     """
     Analyzes the input model and collects the amounts of parameters and MACs
     required to make a forward pass of the model.
@@ -75,6 +77,8 @@ def get_model_complexity_info(model: nn.Module,
     :param output_precision: Floating point precision for representing MACs/params in
         given units.
     :type output_precision: int
+    :param backend_specific_config: Extra configuration for a specific backend.
+    :type backend_specific_config: dict
 
     Returns:
         Tuple[Union[str, int, None], Union[str, int, None]]: Return value is a tuple
@@ -86,14 +90,16 @@ def get_model_complexity_info(model: nn.Module,
     assert isinstance(model, nn.Module)
 
     if FLOPS_BACKEND(backend) == FLOPS_BACKEND.PYTORCH:
-        flops_count, params_count = get_flops_pytorch(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)
+        flops_count, params_count = \
+            get_flops_pytorch(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,
+                              extra_config=backend_specific_config)
     elif FLOPS_BACKEND(backend) == FLOPS_BACKEND.ATEN:
         flops_count, params_count = get_flops_aten(model, input_res,
                                                    print_per_layer_stat,
@@ -102,7 +108,8 @@ def get_model_complexity_info(model: nn.Module,
                                                    custom_modules_hooks,
                                                    output_precision=output_precision,
                                                    flops_units=flops_units,
-                                                   param_units=param_units)
+                                                   param_units=param_units,
+                                                   extra_config=backend_specific_config)
     else:
         raise ValueError('Wrong backend name')
 

ptflops/pytorch_engine.py

@@ -9,7 +9,7 @@
 import sys
 import traceback
 from functools import partial
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
@@ -27,8 +27,9 @@ def get_flops_pytorch(model, input_res,
                       custom_modules_hooks={},
                       output_precision=2,
                       flops_units: Optional[str] = 'GMac',
-                      param_units: Optional[str] = 'M') -> Tuple[Union[int, None],
-                                                                 Union[int, None]]:
+                      param_units: Optional[str] = 'M',
+                      extra_config: Dict = {}) -> Tuple[Union[int, None],
+                                                        Union[int, None]]:
     global CUSTOM_MODULES_MAPPING
     CUSTOM_MODULES_MAPPING = custom_modules_hooks
     flops_model = add_flops_counting_methods(model)
@@ -45,15 +46,18 @@ def get_flops_pytorch(model, input_res,
         except StopIteration:
             batch = torch.ones(()).new_empty((1, *input_res))
 
+    enable_func_ops_patching = extra_config.get('count_functional', True)
     torch_functional_flops = []
     torch_tensor_ops_flops = []
-    patch_functional(torch_functional_flops)
-    patch_tensor_ops(torch_tensor_ops_flops)
+    if enable_func_ops_patching:
+        patch_functional(torch_functional_flops)
+        patch_tensor_ops(torch_tensor_ops_flops)
 
     def reset_environment():
         flops_model.stop_flops_count()
-        unpatch_functional()
-        unpatch_tensor_ops()
+        if enable_func_ops_patching:
+            unpatch_functional()
+            unpatch_tensor_ops()
         global CUSTOM_MODULES_MAPPING
         CUSTOM_MODULES_MAPPING = {}
 

tests/common_test.py

@@ -22,7 +22,6 @@ def forward(self, x):
 
         return CustomModel()
 
-
     @pytest.mark.parametrize("backend", [FLOPS_BACKEND.PYTORCH, FLOPS_BACKEND.ATEN])
     def test_conv(self, default_input_image_size, backend: FLOPS_BACKEND):
         net = nn.Sequential(nn.Conv2d(3, 2, 3, bias=True))
@@ -152,3 +151,14 @@ def test_aten_custom(self, simple_model_mm):
 
         assert params == 0
         assert macs == reference
+
+    def test_torch_ignore_func(self, simple_model_mm):
+        macs, params = \
+            get_model_complexity_info(simple_model_mm, (10, ),
+                                      backend=FLOPS_BACKEND.PYTORCH,
+                                      as_strings=False,
+                                      print_per_layer_stat=False,
+                                      backend_specific_config={'count_functional': False})
+
+        assert params == 0
+        assert macs == 0