Spaces:
Sleeping
Sleeping
| # Copyright (c) OpenMMLab. All rights reserved. | |
| 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), 'input': (3, 16), 'flops': 96.0, 'params': 6.0}, # noqa: E501 | |
| {'model': nn.BatchNorm2d(3), 'input': (3, 16, 16), 'flops': 1536.0, 'params': 6.0}, # noqa: E501 | |
| {'model': nn.BatchNorm3d(3), 'input': (3, 3, 16, 16), 'flops': 4608.0, 'params': 6.0}, # noqa: E501 | |
| {'model': nn.GroupNorm(2, 6), 'input': (6, 16, 16), 'flops': 3072.0, 'params': 12.0}, # noqa: E501 | |
| {'model': nn.InstanceNorm1d(3, affine=True), 'input': (3, 16), 'flops': 96.0, 'params': 6.0}, # noqa: E501 | |
| {'model': nn.InstanceNorm2d(3, affine=True), 'input': (3, 16, 16), 'flops': 1536.0, 'params': 6.0}, # noqa: E501 | |
| {'model': nn.InstanceNorm3d(3, affine=True), 'input': (3, 3, 16, 16), 'flops': 4608.0, 'params': 6.0}, # noqa: E501 | |
| {'model': nn.LayerNorm((3, 16, 16)), 'input': (3, 16, 16), 'flops': 1536.0, 'params': 1536.0}, # noqa: E501 | |
| {'model': nn.LayerNorm((3, 16, 16), elementwise_affine=False), 'input': (3, 16, 16), 'flops': 768.0, 'params': 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' | |