from __future__ import division, print_function, absolute_import
import math
import numpy as np
from itertools import repeat
from collections import namedtuple, defaultdict
import torch
__all__ = ['compute_model_complexity']
"""
Utility
"""
def _ntuple(n):
def parse(x):
if isinstance(x, int):
return tuple(repeat(x, n))
return x
return parse
_single = _ntuple(1)
_pair = _ntuple(2)
_triple = _ntuple(3)
"""
Convolution
"""
def hook_convNd(m, x, y):
k = torch.prod(torch.Tensor(m.kernel_size)).item()
cin = m.in_channels
flops_per_ele = k * cin # + (k*cin-1)
if m.bias is not None:
flops_per_ele += 1
flops = flops_per_ele * y.numel() / m.groups
return int(flops)
"""
Pooling
"""
def hook_maxpool1d(m, x, y):
flops_per_ele = m.kernel_size - 1
flops = flops_per_ele * y.numel()
return int(flops)
def hook_maxpool2d(m, x, y):
k = _pair(m.kernel_size)
k = torch.prod(torch.Tensor(k)).item()
# ops: compare
flops_per_ele = k - 1
flops = flops_per_ele * y.numel()
return int(flops)
def hook_maxpool3d(m, x, y):
k = _triple(m.kernel_size)
k = torch.prod(torch.Tensor(k)).item()
flops_per_ele = k - 1
flops = flops_per_ele * y.numel()
return int(flops)
def hook_avgpool1d(m, x, y):
flops_per_ele = m.kernel_size
flops = flops_per_ele * y.numel()
return int(flops)
def hook_avgpool2d(m, x, y):
k = _pair(m.kernel_size)
k = torch.prod(torch.Tensor(k)).item()
flops_per_ele = k
flops = flops_per_ele * y.numel()
return int(flops)
def hook_avgpool3d(m, x, y):
k = _triple(m.kernel_size)
k = torch.prod(torch.Tensor(k)).item()
flops_per_ele = k
flops = flops_per_ele * y.numel()
return int(flops)
def hook_adapmaxpool1d(m, x, y):
x = x[0]
out_size = m.output_size
k = math.ceil(x.size(2) / out_size)
flops_per_ele = k - 1
flops = flops_per_ele * y.numel()
return int(flops)
def hook_adapmaxpool2d(m, x, y):
x = x[0]
out_size = _pair(m.output_size)
k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)
k = torch.prod(torch.ceil(k)).item()
flops_per_ele = k - 1
flops = flops_per_ele * y.numel()
return int(flops)
def hook_adapmaxpool3d(m, x, y):
x = x[0]
out_size = _triple(m.output_size)
k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)
k = torch.prod(torch.ceil(k)).item()
flops_per_ele = k - 1
flops = flops_per_ele * y.numel()
return int(flops)
def hook_adapavgpool1d(m, x, y):
x = x[0]
out_size = m.output_size
k = math.ceil(x.size(2) / out_size)
flops_per_ele = k
flops = flops_per_ele * y.numel()
return int(flops)
def hook_adapavgpool2d(m, x, y):
x = x[0]
out_size = _pair(m.output_size)
k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)
k = torch.prod(torch.ceil(k)).item()
flops_per_ele = k
flops = flops_per_ele * y.numel()
return int(flops)
def hook_adapavgpool3d(m, x, y):
x = x[0]
out_size = _triple(m.output_size)
k = torch.Tensor(list(x.size()[2:])) / torch.Tensor(out_size)
k = torch.prod(torch.ceil(k)).item()
flops_per_ele = k
flops = flops_per_ele * y.numel()
return int(flops)
"""
Non-linear activations
"""
def hook_relu(m, x, y):
# eq: max(0, x)
num_ele = y.numel()
return int(num_ele)
def hook_leakyrelu(m, x, y):
# eq: max(0, x) + negative_slope*min(0, x)
num_ele = y.numel()
flops = 3 * num_ele
return int(flops)
"""
Normalization
"""
def hook_batchnormNd(m, x, y):
num_ele = y.numel()
flops = 2 * num_ele # mean and std
if m.affine:
flops += 2 * num_ele # gamma and beta
return int(flops)
def hook_instancenormNd(m, x, y):
return hook_batchnormNd(m, x, y)
def hook_groupnorm(m, x, y):
return hook_batchnormNd(m, x, y)
def hook_layernorm(m, x, y):
num_ele = y.numel()
flops = 2 * num_ele # mean and std
if m.elementwise_affine:
flops += 2 * num_ele # gamma and beta
return int(flops)
"""
Linear
"""
def hook_linear(m, x, y):
flops_per_ele = m.in_features # + (m.in_features-1)
if m.bias is not None:
flops_per_ele += 1
flops = flops_per_ele * y.numel()
return int(flops)
__generic_flops_counter = {
# Convolution
'Conv1d': hook_convNd,
'Conv2d': hook_convNd,
'Conv3d': hook_convNd,
# Pooling
'MaxPool1d': hook_maxpool1d,
'MaxPool2d': hook_maxpool2d,
'MaxPool3d': hook_maxpool3d,
'AvgPool1d': hook_avgpool1d,
'AvgPool2d': hook_avgpool2d,
'AvgPool3d': hook_avgpool3d,
'AdaptiveMaxPool1d': hook_adapmaxpool1d,
'AdaptiveMaxPool2d': hook_adapmaxpool2d,
'AdaptiveMaxPool3d': hook_adapmaxpool3d,
'AdaptiveAvgPool1d': hook_adapavgpool1d,
'AdaptiveAvgPool2d': hook_adapavgpool2d,
'AdaptiveAvgPool3d': hook_adapavgpool3d,
# Non-linear activations
'ReLU': hook_relu,
'ReLU6': hook_relu,
'LeakyReLU': hook_leakyrelu,
# Normalization
'BatchNorm1d': hook_batchnormNd,
'BatchNorm2d': hook_batchnormNd,
'BatchNorm3d': hook_batchnormNd,
'InstanceNorm1d': hook_instancenormNd,
'InstanceNorm2d': hook_instancenormNd,
'InstanceNorm3d': hook_instancenormNd,
'GroupNorm': hook_groupnorm,
'LayerNorm': hook_layernorm,
# Linear
'Linear': hook_linear,
}
__conv_linear_flops_counter = {
# Convolution
'Conv1d': hook_convNd,
'Conv2d': hook_convNd,
'Conv3d': hook_convNd,
# Linear
'Linear': hook_linear,
}
def _get_flops_counter(only_conv_linear):
if only_conv_linear:
return __conv_linear_flops_counter
return __generic_flops_counter
[docs]def compute_model_complexity(
model, input_size, verbose=False, only_conv_linear=True
):
"""Returns number of parameters and FLOPs.
.. note::
(1) this function only provides an estimate of the theoretical time complexity
rather than the actual running time which depends on implementations and hardware,
and (2) the FLOPs is only counted for layers that are used at test time. This means
that redundant layers such as person ID classification layer will be ignored as it
is discarded when doing feature extraction. Note that the inference graph depends on
how you construct the computations in ``forward()``.
Args:
model (nn.Module): network model.
input_size (tuple): input size, e.g. (1, 3, 256, 128).
verbose (bool, optional): shows detailed complexity of
each module. Default is False.
only_conv_linear (bool, optional): only considers convolution
and linear layers when counting flops. Default is True.
If set to False, flops of all layers will be counted.
Examples::
>>> from torchreid import models, utils
>>> model = models.build_model(name='resnet50', num_classes=1000)
>>> num_params, flops = utils.compute_model_complexity(model, (1, 3, 256, 128), verbose=True)
"""
registered_handles = []
layer_list = []
layer = namedtuple('layer', ['class_name', 'params', 'flops'])
def _add_hooks(m):
def _has_submodule(m):
return len(list(m.children())) > 0
def _hook(m, x, y):
params = sum(p.numel() for p in m.parameters())
class_name = str(m.__class__.__name__)
flops_counter = _get_flops_counter(only_conv_linear)
if class_name in flops_counter:
flops = flops_counter[class_name](m, x, y)
else:
flops = 0
layer_list.append(
layer(class_name=class_name, params=params, flops=flops)
)
# only consider the very basic nn layer
if _has_submodule(m):
return
handle = m.register_forward_hook(_hook)
registered_handles.append(handle)
default_train_mode = model.training
model.eval().apply(_add_hooks)
input = torch.rand(input_size)
if next(model.parameters()).is_cuda:
input = input.cuda()
model(input) # forward
for handle in registered_handles:
handle.remove()
model.train(default_train_mode)
if verbose:
per_module_params = defaultdict(list)
per_module_flops = defaultdict(list)
total_params, total_flops = 0, 0
for layer in layer_list:
total_params += layer.params
total_flops += layer.flops
if verbose:
per_module_params[layer.class_name].append(layer.params)
per_module_flops[layer.class_name].append(layer.flops)
if verbose:
num_udscore = 55
print(' {}'.format('-' * num_udscore))
print(' Model complexity with input size {}'.format(input_size))
print(' {}'.format('-' * num_udscore))
for class_name in per_module_params:
params = int(np.sum(per_module_params[class_name]))
flops = int(np.sum(per_module_flops[class_name]))
print(
' {} (params={:,}, flops={:,})'.format(
class_name, params, flops
)
)
print(' {}'.format('-' * num_udscore))
print(
' Total (params={:,}, flops={:,})'.format(
total_params, total_flops
)
)
print(' {}'.format('-' * num_udscore))
return total_params, total_flops