[Enhance] Use fvcore to calculate FLOPS. (#1000)

* [Feature] Use fvcore for flops count

* update requirements

* update

requirements/optional.txt

@@ -1,3 +1,4 @@
 albumentations>=0.3.2 --no-binary qudida,albumentations
 colorama
+fvcore
 requests

tools/analysis_tools/get_flops.py

@@ -1,8 +1,17 @@
 # Copyright (c) OpenMMLab. All rights reserved.
 import argparse
 
-from mmcv import Config
+import torch
-from mmcv.cnn.utils import get_model_complexity_info
+
+try:
+    from fvcore.nn import (ActivationCountAnalysis, FlopCountAnalysis,
+                           flop_count_str, flop_count_table, parameter_count)
+except ImportError:
+    print('You may need to install fvcore for flops computation, '
+          'and you can use `pip install -r requirements/optional.txt` '
+          'to set up the environment')
+from fvcore.nn.print_model_statistics import _format_size
+from mmengine import Config
 
 from mmcls.models import build_classifier
 
@@ -42,10 +51,28 @@ def main():
             'FLOPs counter is currently not currently supported with {}'.
             format(model.__class__.__name__))
 
-    flops, params = get_model_complexity_info(model, input_shape)
+    inputs = (torch.randn((1, *input_shape)), )
+    flops_ = FlopCountAnalysis(model, inputs)
+    activations_ = ActivationCountAnalysis(model, inputs)
+
+    flops = _format_size(flops_.total())
+    activations = _format_size(activations_.total())
+    params = _format_size(parameter_count(model)[''])
+
+    flop_table = flop_count_table(
+        flops=flops_,
+        activations=activations_,
+        show_param_shapes=True,
+    )
+    flop_str = flop_count_str(flops=flops_, activations=activations_)
+
+    print('\n' + flop_str)
+    print('\n' + flop_table)
+
     split_line = '=' * 30
     print(f'{split_line}\nInput shape: {input_shape}\n'
-          f'Flops: {flops}\nParams: {params}\n{split_line}')
+          f'Flops: {flops}\nParams: {params}\n'
+          f'Activation: {activations}\n{split_line}')
     print('!!!Please be cautious if you use the results in papers. '
           'You may need to check if all ops are supported and verify that the '
           'flops computation is correct.')