File size: 2,105 Bytes
b4f7b8c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 |
'''
Copyright (C) 2019-2021 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 torch.nn as nn
from .pytorch_engine import get_flops_pytorch
from .utils import flops_to_string, params_to_string
def get_model_complexity_info(model, input_res,
print_per_layer_stat=True,
as_strings=True,
input_constructor=None, ost=sys.stdout,
verbose=False, ignore_modules=[],
custom_modules_hooks={}, backend='pytorch',
flops_units=None, param_units=None,
output_precision=2):
assert type(input_res) is tuple
assert len(input_res) >= 1
assert isinstance(model, nn.Module)
if 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)
else:
raise ValueError('Wrong backend name')
if as_strings:
flops_string = flops_to_string(
flops_count,
units=flops_units,
precision=output_precision
)
params_string = params_to_string(
params_count,
units=param_units,
precision=output_precision
)
return flops_string, params_string
return flops_count, params_count
|