Add utils to calculate model complexity and adjust folder structure for mmcv/cnn/utils . (#366)

* Add utils to calculate model complexity info

* remove _InstanceNorm in unittest

* add docstring and increase unittest coverage

* fix deconv_flops_counter_hook to accept different data shape

* test when model is not a common instance

* put flops_counter.py and weight_init.py into mmcv/cnn/utils folder

* fix import name

* reformat some docstrings

* update the scripts with latest one and remove redundant codes

* directly represent a model without string and eval()

* reformat code

mmcv/cnn/__init__.py

@@ -6,10 +6,10 @@ from .bricks import (ACTIVATION_LAYERS, CONV_LAYERS, NORM_LAYERS,
                      build_conv_layer, build_norm_layer, build_padding_layer,
                      build_upsample_layer, is_norm)
 from .resnet import ResNet, make_res_layer
+from .utils import (bias_init_with_prob, caffe2_xavier_init, constant_init,
+                    get_model_complexity_info, kaiming_init, normal_init,
+                    uniform_init, xavier_init)
 from .vgg import VGG, make_vgg_layer
-from .weight_init import (bias_init_with_prob, caffe2_xavier_init,
-                          constant_init, kaiming_init, normal_init,
-                          uniform_init, xavier_init)
 
 __all__ = [
     'AlexNet', 'VGG', 'make_vgg_layer', 'ResNet', 'make_res_layer',
@@ -18,5 +18,6 @@ __all__ = [
     'build_activation_layer', 'build_conv_layer', 'build_norm_layer',
     'build_padding_layer', 'build_upsample_layer', 'is_norm', 'NonLocal1d',
     'NonLocal2d', 'NonLocal3d', 'ACTIVATION_LAYERS', 'CONV_LAYERS',
-    'NORM_LAYERS', 'PADDING_LAYERS', 'UPSAMPLE_LAYERS', 'Scale'
+    'NORM_LAYERS', 'PADDING_LAYERS', 'UPSAMPLE_LAYERS', 'Scale',
+    'get_model_complexity_info'
 ]

mmcv/cnn/bricks/conv_module.py

@@ -2,7 +2,7 @@ import warnings
 
 import torch.nn as nn
 
-from ..weight_init import constant_init, kaiming_init
+from ..utils import constant_init, kaiming_init
 from .activation import build_activation_layer
 from .conv import build_conv_layer
 from .norm import build_norm_layer

mmcv/cnn/bricks/non_local.py

@@ -3,7 +3,7 @@ from abc import ABCMeta
 import torch
 import torch.nn as nn
 
-from ..weight_init import constant_init, normal_init
+from ..utils import constant_init, normal_init
 from .conv_module import ConvModule
 
 

mmcv/cnn/bricks/upsample.py

 import torch.nn as nn
 import torch.nn.functional as F
 
-from ..weight_init import xavier_init
+from ..utils import xavier_init
 from .registry import UPSAMPLE_LAYERS
 
 UPSAMPLE_LAYERS.register_module('nearest', module=nn.Upsample)

mmcv/cnn/resnet.py

@@ -5,7 +5,7 @@ import torch.nn as nn
 import torch.utils.checkpoint as cp
 
 from ..runner import load_checkpoint
-from .weight_init import constant_init, kaiming_init
+from .utils import constant_init, kaiming_init
 
 
 def conv3x3(in_planes, out_planes, stride=1, dilation=1):

mmcv/cnn/utils/__init__.py

+# Copyright (c) Open-MMLab. All rights reserved.
+from .flops_counter import get_model_complexity_info
+from .weight_init import (bias_init_with_prob, caffe2_xavier_init,
+                          constant_init, kaiming_init, normal_init,
+                          uniform_init, xavier_init)
+
+__all__ = [
+    'get_model_complexity_info', 'bias_init_with_prob', 'caffe2_xavier_init',
+    'constant_init', 'kaiming_init', 'normal_init', 'uniform_init',
+    'xavier_init'
+]

mmcv/cnn/vgg.py

@@ -4,7 +4,7 @@ import logging
 import torch.nn as nn
 
 from ..runner import load_checkpoint
-from .weight_init import constant_init, kaiming_init, normal_init
+from .utils import constant_init, kaiming_init, normal_init
 
 
 def conv3x3(in_planes, out_planes, dilation=1):

tests/test_cnn/test_flops_counter.py

+import pytest
+import torch
+import torch.nn as nn
+
+from mmcv.cnn import get_model_complexity_info
+from mmcv.cnn.utils.flops_counter import flops_to_string, params_to_string
+
+try:
+    from StringIO import StringIO
+except ImportError:
+    from io import StringIO
+
+# yapf: disable
+gt_results = [
+    {'model': nn.Conv1d(3, 8, 3), 'input': (3, 16), 'flops': 1120.0, 'params': 80.0},  # noqa: E501
+    {'model': nn.Conv2d(3, 8, 3), 'input': (3, 16, 16), 'flops': 43904.0, 'params': 224.0},  # noqa: E501
+    {'model': nn.Conv3d(3, 8, 3), 'input': (3, 3, 16, 16), 'flops': 128576.0, 'params': 656.0},  # noqa: E501
+    {'model': nn.ReLU(), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.PReLU(), 'input': (3, 16, 16), 'flops': 768.0, 'params': 1},  # noqa: E501
+    {'model': nn.ELU(), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.LeakyReLU(), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.ReLU6(), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.MaxPool1d(2), 'input': (3, 16), 'flops': 48.0, 'params': 0},  # noqa: E501
+    {'model': nn.MaxPool2d(2), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.MaxPool3d(2), 'input': (3, 3, 16, 16), 'flops': 2304.0, 'params': 0},  # noqa: E501
+    {'model': nn.AvgPool1d(2), 'input': (3, 16), 'flops': 48.0, 'params': 0},  # noqa: E501
+    {'model': nn.AvgPool2d(2), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.AvgPool3d(2), 'input': (3, 3, 16, 16), 'flops': 2304.0, 'params': 0},  # noqa: E501
+    {'model': nn.AdaptiveMaxPool1d(2), 'input': (3, 16), 'flops': 48.0, 'params': 0},  # noqa: E501
+    {'model': nn.AdaptiveMaxPool2d(2), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.AdaptiveMaxPool3d(2), 'input': (3, 3, 16, 16), 'flops': 2304.0, 'params': 0},  # noqa: E501
+    {'model': nn.AdaptiveAvgPool1d(2), 'input': (3, 16), 'flops': 48.0, 'params': 0},  # noqa: E501
+    {'model': nn.AdaptiveAvgPool2d(2), 'input': (3, 16, 16), 'flops': 768.0, 'params': 0},  # noqa: E501
+    {'model': nn.AdaptiveAvgPool3d(2), 'input': (3, 3, 16, 16), 'flops': 2304.0, 'params': 0},  # noqa: E501
+    {'model': nn.BatchNorm1d(3, 8), 'input': (3, 16), 'flops': 96.0, 'params': 6.0},  # noqa: E501
+    {'model': nn.BatchNorm2d(3, 8), 'input': (3, 16, 16), 'flops': 1536.0, 'params': 6.0},  # noqa: E501
+    {'model': nn.BatchNorm3d(3, 8), 'input': (3, 3, 16, 16), 'flops': 4608.0, 'params': 6.0},  # noqa: E501
+    {'model': nn.Linear(1024, 2), 'input': (1024, ), 'flops': 2048.0, 'params': 2050.0},  # noqa: E501
+    {'model': nn.ConvTranspose2d(3, 8, 3), 'input': (3, 16, 16), 'flops': 57888, 'params': 224.0},  # noqa: E501
+    {'model': nn.Upsample((32, 32)), 'input': (3, 16, 16), 'flops': 3072.0, 'params': 0}  # noqa: E501
+]
+# yapf: enable
+
+
+class ExampleModel(nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv2d = nn.Conv2d(3, 8, 3)
+
+    def forward(self, imgs):
+        x = torch.randn((1, *imgs))
+        return self.conv2d(x)
+
+
+def input_constructor(x):
+    return dict(imgs=x)
+
+
+def test_flops_counter():
+    with pytest.raises(AssertionError):
+        # input_res should be a tuple
+        model = nn.Conv2d(3, 8, 3)
+        input_res = [1, 3, 16, 16]
+        get_model_complexity_info(model, input_res)
+
+    with pytest.raises(AssertionError):
+        # len(input_res) >= 2
+        model = nn.Conv2d(3, 8, 3)
+        input_res = tuple()
+        get_model_complexity_info(model, input_res)
+
+    # test common layers
+    for item in gt_results:
+        model = item['model']
+        input = item['input']
+        flops, params = get_model_complexity_info(
+            model, input, as_strings=False, print_per_layer_stat=False)
+        assert flops == item['flops'] and params == item['params']
+
+    # test input constructor
+    model = ExampleModel()
+    x = (3, 16, 16)
+    flops, params = get_model_complexity_info(
+        model,
+        x,
+        as_strings=False,
+        print_per_layer_stat=False,
+        input_constructor=input_constructor)
+    assert flops == 43904.0 and params == 224.0
+
+    # test output string
+    model = nn.Conv3d(3, 8, 3)
+    x = (3, 3, 512, 512)
+    flops, params = get_model_complexity_info(
+        model, x, print_per_layer_stat=False)
+    assert flops == '0.17 GFLOPs' and params == str(656)
+
+    # test print per layer status
+    model = nn.Conv1d(3, 8, 3)
+    x = (3, 16)
+    out = StringIO()
+    get_model_complexity_info(model, x, ost=out)
+    assert out.getvalue() == \
+        'Conv1d(0.0 M, 100.000% Params, 0.0 GFLOPs, 100.000% FLOPs, 3, 8, kernel_size=(3,), stride=(1,))\n'  # noqa: E501
+
+    # test when model is not a common instance
+    model = nn.Sequential(nn.Conv2d(3, 8, 3), nn.Flatten(), nn.Linear(1568, 2))
+    x = (3, 16, 16)
+    flops, params = get_model_complexity_info(
+        model, x, as_strings=False, print_per_layer_stat=True)
+    assert flops == 47040.0 and params == 3362
+
+
+def test_flops_to_string():
+    flops = 6.54321 * 10.**9
+    assert flops_to_string(flops) == '6.54 GFLOPs'
+    assert flops_to_string(flops, 'MFLOPs') == '6543.21 MFLOPs'
+    assert flops_to_string(flops, 'KFLOPs') == '6543210.0 KFLOPs'
+    assert flops_to_string(flops, 'FLOPs') == '6543210000.0 FLOPs'
+    assert flops_to_string(flops, precision=4) == '6.5432 GFLOPs'
+
+    flops = 6.54321 * 10.**9
+    assert flops_to_string(flops, None) == '6.54 GFLOPs'
+    flops = 3.21 * 10.**7
+    assert flops_to_string(flops, None) == '32.1 MFLOPs'
+    flops = 5.4 * 10.**3
+    assert flops_to_string(flops, None) == '5.4 KFLOPs'
+    flops = 987
+    assert flops_to_string(flops, None) == '987 FLOPs'
+
+
+def test_params_to_string():
+    num_params = 3.21 * 10.**7
+    assert params_to_string(num_params) == '32.1 M'
+    num_params = 4.56 * 10.**5
+    assert params_to_string(num_params) == '456.0 k'
+    num_params = 7.89 * 10.**2
+    assert params_to_string(num_params) == '789.0'
+
+    num_params = 6.54321 * 10.**7
+    assert params_to_string(num_params, 'M') == '65.43 M'
+    assert params_to_string(num_params, 'K') == '65432.1 K'
+    assert params_to_string(num_params, '') == '65432100.0'
+    assert params_to_string(num_params, precision=4) == '65.4321 M'