# MegEngine is Licensed under the Apache License, Version 2.0 (the "License")
#
# Copyright (c) 2014-2021 Megvii Inc. All rights reserved.
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT ARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
from collections import Iterable, namedtuple
from functools import partial
from typing import Iterable
import numpy as np
import tabulate
import megengine as mge
import megengine.module as m
import megengine.module.qat as qatm
import megengine.module.quantized as qm
from megengine import Tensor
from megengine import functional as F
from megengine.core.tensor.dtype import get_dtype_bit
from megengine.functional.tensor import zeros
from megengine.tensor import Tensor
from .module_utils import set_module_mode_safe
try:
mge.logger.MegEngineLogFormatter.max_lines = float("inf")
except AttributeError as e:
raise ValueError("set logger max lines failed")
logger = mge.get_logger(__name__)
logger.setLevel("INFO")
_calc_flops_dict = {}
_calc_receptive_field_dict = {}
def _receptive_field_fallback(module, inputs, outputs):
if not _receptive_field_enabled:
return
assert not hasattr(module, "_rf")
assert not hasattr(module, "_stride")
if len(inputs) == 0:
# TODO: support other dimension
module._rf = (1, 1)
module._stride = (1, 1)
return module._rf, module._stride
rf, stride = preprocess_receptive_field(module, inputs, outputs)
module._rf = rf
module._stride = stride
return rf, stride
# key tuple, impl_dict, fallback
_iter_list = [
("flops_num", _calc_flops_dict, None),
(
("receptive_field", "stride"),
_calc_receptive_field_dict,
_receptive_field_fallback,
),
]
_receptive_field_enabled = False
def _register_dict(*modules, dict=None):
def callback(impl):
for module in modules:
dict[module] = impl
return impl
return callback
def register_flops(*modules):
return _register_dict(*modules, dict=_calc_flops_dict)
def register_receptive_field(*modules):
return _register_dict(*modules, dict=_calc_receptive_field_dict)
def enable_receptive_field():
global _receptive_field_enabled
_receptive_field_enabled = True
def disable_receptive_field():
global _receptive_field_enabled
_receptive_field_enabled = False
@register_flops(
m.Conv1d, m.Conv2d, m.Conv3d, m.ConvTranspose2d, m.LocalConv2d, m.DeformableConv2d
)
def flops_convNd(module: m.Conv2d, inputs, outputs):
bias = 1 if module.bias is not None else 0
# N x Cout x H x W x (Cin x Kw x Kh + bias)
return np.prod(outputs[0].shape) * (
module.in_channels // module.groups * np.prod(module.kernel_size) + bias
)
@register_flops(
m.batchnorm._BatchNorm, m.SyncBatchNorm, m.GroupNorm, m.LayerNorm, m.InstanceNorm,
)
def flops_norm(module: m.Linear, inputs, outputs):
return np.prod(inputs[0].shape) * 7
@register_flops(m.AvgPool2d, m.MaxPool2d)
def flops_pool(module: m.AvgPool2d, inputs, outputs):
kernel_sum = 0
if isinstance(module.kernel_size, tuple) and len(module.kernel_size) == 2:
kernel_sum = np.prod(module.kernel_size)
else:
kernel_sum = module.kernel_size ** 2
return np.prod(outputs[0].shape) * kernel_sum
@register_flops(m.AdaptiveAvgPool2d, m.AdaptiveMaxPool2d)
def flops_adaptivePool(module: m.AdaptiveAvgPool2d, inputs, outputs):
stride_h = np.floor(inputs[0].shape[2] / (inputs[0].shape[2] - 1))
kernel_h = inputs[0].shape[2] - (inputs[0].shape[2] - 1) * stride_h
stride_w = np.floor(inputs[0].shape[3] / (inputs[0].shape[3] - 1))
kernel_w = inputs[0].shape[3] - (inputs[0].shape[3] - 1) * stride_w
return np.prod(outputs[0].shape) * kernel_h * kernel_w
@register_flops(m.Linear)
def flops_linear(module: m.Linear, inputs, outputs):
bias = module.out_features if module.bias is not None else 0
return np.prod(outputs[0].shape) * module.in_features + bias
@register_flops(m.BatchMatMulActivation)
def flops_batchmatmul(module: m.BatchMatMulActivation, inputs, outputs):
bias = 1 if module.bias is not None else 0
x = inputs[0]
w = module.weight
batch_size = x.shape[0]
n, p = x.shape[1:]
_, m = w.shape[1:]
return n * (p + bias) * m * batch_size
# does not need import qat and quantized module since they inherit from float module.
hook_modules = (
m.conv._ConvNd,
m.Linear,
m.BatchMatMulActivation,
m.batchnorm._BatchNorm,
m.LayerNorm,
m.GroupNorm,
m.InstanceNorm,
m.pooling._PoolNd,
m.adaptive_pooling._AdaptivePoolNd,
)
def _mean(inp):
inp = mge.tensor(inp).astype(np.float32)
return F.mean(inp).numpy()
def _std(inp):
inp = mge.tensor(inp).astype(np.float32)
return F.std(inp).numpy()
def dict2table(list_of_dict, header):
table_data = [header]
for d in list_of_dict:
row = []
for h in header:
v = ""
if h in d:
v = d[h]
row.append(v)
table_data.append(row)
return table_data
def sizeof_fmt(num, suffix="B"):
if suffix == "B":
scale = 1024.0
units = ["", "Ki", "Mi", "Gi", "Ti", "Pi", "Ei", "Zi", "Yi"]
else:
scale = 1000.0
units = ["", "K", "M", "G", "T", "P", "E", "Z", "Y"]
for unit in units:
if abs(num) < scale or unit == units[-1]:
return "{:3.3f} {}{}".format(num, unit, suffix)
num /= scale
def preprocess_receptive_field(module, inputs, outputs):
# TODO: support other dimensions
pre_rf = (
max(getattr(i.owner, "_rf", (1, 1))[0] for i in inputs),
max(getattr(i.owner, "_rf", (1, 1))[1] for i in inputs),
)
pre_stride = (
max(getattr(i.owner, "_stride", (1, 1))[0] for i in inputs),
max(getattr(i.owner, "_stride", (1, 1))[1] for i in inputs),
)
return pre_rf, pre_stride
def get_op_stats(module, inputs, outputs):
if not isinstance(outputs, tuple) and not isinstance(outputs, list):
outputs = (outputs,)
rst = {
"input_shapes": [i.shape for i in inputs],
"output_shapes": [o.shape for o in outputs],
}
valid_flag = False
for key, _dict, fallback in _iter_list:
for _type in _dict:
if isinstance(module, _type):
value = _dict[_type](module, inputs, outputs)
valid_flag = True
break
else:
if fallback is not None:
value = fallback(module, inputs, outputs)
continue
if isinstance(key, tuple):
assert isinstance(value, tuple)
for k, v in zip(key, value):
rst[k] = v
else:
rst[key] = value
if valid_flag:
return rst
else:
return None
return
def sum_op_stats(flops, bar_length_max=20):
max_flops_num = max([i["flops_num"] for i in flops] + [0])
total_flops_num = 0
for d in flops:
total_flops_num += int(d["flops_num"])
d["flops_cum"] = sizeof_fmt(total_flops_num, suffix="OPs")
for d in flops:
ratio = d["ratio"] = d["flops_num"] / total_flops_num
d["percentage"] = "{:.2f}%".format(ratio * 100)
bar_length = int(d["flops_num"] / max_flops_num * bar_length_max)
d["bar"] = "#" * bar_length
d["flops"] = sizeof_fmt(d["flops_num"], suffix="OPs")
total_flops_str = sizeof_fmt(total_flops_num, suffix="OPs")
total_var_size = sum(
sum(s[1] if len(s) > 1 else 0 for s in d["output_shapes"]) for d in flops
)
flops.append(
dict(name="total", flops=total_flops_str, output_shapes=total_var_size)
)
return total_flops_num, flops
def print_op_stats(flops):
header = [
"name",
"class_name",
"input_shapes",
"output_shapes",
"flops",
"flops_cum",
"percentage",
"bar",
]
if _receptive_field_enabled:
header.insert(4, "receptive_field")
header.insert(5, "stride")
logger.info("flops stats: \n" + tabulate.tabulate(dict2table(flops, header=header)))
def get_param_stats(param: Tensor):
nbits = get_dtype_bit(np.dtype(param.dtype).name)
shape = param.shape
param_dim = np.prod(param.shape)
param_size = param_dim * nbits // 8
return {
"dtype": np.dtype(param.dtype),
"shape": shape,
"mean": "{:.3g}".format(_mean(param)),
"std": "{:.3g}".format(_std(param)),
"param_dim": param_dim,
"nbits": nbits,
"size": param_size,
}
def sum_param_stats(params, bar_length_max=20):
max_size = max([d["size"] for d in params] + [0])
total_param_dims, total_param_size = 0, 0
for d in params:
total_param_dims += int(d["param_dim"])
total_param_size += int(d["size"])
d["size_cum"] = sizeof_fmt(total_param_size)
for d in params:
ratio = d["size"] / total_param_size
d["ratio"] = ratio
d["percentage"] = "{:.2f}%".format(ratio * 100)
bar_length = int(d["size"] / max_size * bar_length_max)
d["size_bar"] = "#" * bar_length
d["size"] = sizeof_fmt(d["size"])
param_size = sizeof_fmt(total_param_size)
params.append(dict(name="total", param_dim=total_param_dims, size=param_size,))
return total_param_dims, total_param_size, params
def print_param_stats(params):
header = [
"name",
"dtype",
"shape",
"mean",
"std",
"param_dim",
"nbits",
"size",
"size_cum",
"percentage",
"size_bar",
]
logger.info(
"param stats: \n" + tabulate.tabulate(dict2table(params, header=header))
)
def get_activation_stats(output: Tensor, has_input=False):
out_shape = output.shape
activations_dtype = np.dtype(output.dtype)
nbits = get_dtype_bit(activations_dtype.name)
act_dim = np.prod(out_shape)
act_size = act_dim * nbits // 8
activation_stats = {
"dtype": activations_dtype,
"shape": out_shape,
"act_dim": act_dim,
"nbits": nbits,
"size": act_size,
}
if has_input:
activation_stats["mean"] = "{:.3g}".format(_mean(output))
activation_stats["std"] = "{:.3g}".format(_std(output))
return activation_stats
def sum_activations_stats(activations, bar_length_max=20):
max_act_size = max([i["size"] for i in activations] + [0])
total_act_dims, total_act_size = 0, 0
for d in activations:
total_act_size += int(d["size"])
total_act_dims += int(d["act_dim"])
d["size_cum"] = sizeof_fmt(total_act_size)
for d in activations:
ratio = d["ratio"] = d["size"] / total_act_size
d["percentage"] = "{:.2f}%".format(ratio * 100)
bar_length = int(d["size"] / max_act_size * bar_length_max)
d["size_bar"] = "#" * bar_length
d["size"] = sizeof_fmt(d["size"])
act_size = sizeof_fmt(total_act_size)
activations.append(dict(name="total", act_dim=total_act_dims, size=act_size,))
return total_act_dims, total_act_size, activations
def print_activations_stats(activations, has_input=False):
header = [
"name",
"class_name",
"dtype",
"shape",
"nbits",
"act_dim",
"size",
"size_cum",
"percentage",
"size_bar",
]
if has_input:
header.insert(4, "mean")
header.insert(5, "std")
logger.info(
"activations stats: \n"
+ tabulate.tabulate(dict2table(activations, header=header))
)
def print_summary(**kwargs):
data = [["item", "value"]]
data.extend(list(kwargs.items()))
logger.info("summary\n" + tabulate.tabulate(data))
[文档]def module_stats(
model: m.Module,
inputs: Iterable[np.ndarray] = None,
input_shapes: list = None,
cal_params: bool = True,
cal_flops: bool = True,
cal_activations: bool = True,
logging_to_stdout: bool = True,
bar_length_max: int = 20,
):
r"""Calculate and print ``model``'s statistics by adding hook and record Module's inputs outputs size.
Args:
model: model that need to get stats info.
inputs: user defined input data for running model and calculating stats, alternative with input_shapes.
input_shapes: shapes to generate random inputs for running model and calculating stats, alternative with inputs.
cal_params: whether calculate and record params size.
cal_flops: whether calculate and record op flops.
cal_activations: whether calculate and record op activations.
logging_to_stdout: whether print all calculated statistic details.
bar_length_max: size of bar indicating max flops or parameter size in net stats.
"""
has_inputs = False
if inputs is not None:
has_inputs = True
if not isinstance(inputs, (tuple, list)):
inputs = [inputs]
inputs = [Tensor(input, dtype=np.float32) for input in inputs]
else:
if input_shapes:
if not isinstance(input_shapes[0], tuple):
input_shapes = [input_shapes]
inputs = [zeros(in_size, dtype=np.float32) for in_size in input_shapes]
else:
logger.error(
"Inputs or input_shapes is required for running model and calculating stats.",
exc_info=True,
)
return
if not cal_activations:
log_activations = False
disable_receptive_field()
def module_stats_hook(module, inputs, outputs, name=""):
class_name = str(module.__class__).split(".")[-1].split("'")[0]
if cal_flops:
flops_stats = get_op_stats(module, inputs, outputs)
if flops_stats is not None:
flops_stats["name"] = name
flops_stats["class_name"] = class_name
flops.append(flops_stats)
if cal_params:
if hasattr(module, "weight") and module.weight is not None:
w = module.weight
param_stats = get_param_stats(w)
param_stats["name"] = name + "-w"
params.append(param_stats)
if hasattr(module, "bias") and module.bias is not None:
b = module.bias
param_stats = get_param_stats(b)
param_stats["name"] = name + "-b"
params.append(param_stats)
if cal_activations:
if not isinstance(outputs, (tuple, list)):
output = outputs
else:
output = outputs[0]
activation_stats = get_activation_stats(output, has_inputs)
activation_stats["name"] = name
activation_stats["class_name"] = class_name
activations.append(activation_stats)
params = []
flops = []
hooks = []
activations = []
total_stats = namedtuple(
"total_stats", ["param_size", "param_dims", "flops", "act_size", "act_dims"]
)
stats_details = namedtuple("module_stats", ["params", "flops", "activations"])
for (name, module) in model.named_modules():
if isinstance(module, hook_modules):
hooks.append(
module.register_forward_hook(partial(module_stats_hook, name=name))
)
with set_module_mode_safe(model, training=False) as model:
model(*inputs)
for h in hooks:
h.remove()
extra_info = {
"#params": len(params),
}
(
total_flops,
total_param_dims,
total_param_size,
total_act_dims,
total_act_size,
) = (0, 0, 0, 0, 0)
if cal_params:
total_param_dims, total_param_size, params = sum_param_stats(
params, bar_length_max
)
extra_info["total_param_dims"] = sizeof_fmt(total_param_dims, suffix="")
extra_info["total_param_size"] = sizeof_fmt(total_param_size)
if logging_to_stdout:
print_param_stats(params)
if cal_flops:
total_flops, flops = sum_op_stats(flops, bar_length_max)
extra_info["total_flops"] = sizeof_fmt(total_flops, suffix="OPs")
if logging_to_stdout:
print_op_stats(flops)
if cal_activations:
total_act_dims, total_act_size, activations = sum_activations_stats(
activations, bar_length_max
)
extra_info["total_act_dims"] = sizeof_fmt(total_act_dims, suffix="")
extra_info["total_act_size"] = sizeof_fmt(total_act_size)
if logging_to_stdout:
print_activations_stats(activations, has_inputs)
if cal_flops and cal_params and total_param_size != 0:
extra_info["flops/param_size"] = "{:3.3f}".format(
total_flops / total_param_size
)
print_summary(**extra_info)
return (
total_stats(
param_size=total_param_size,
param_dims=total_param_dims,
flops=total_flops,
act_size=total_act_size,
act_dims=total_act_dims,
),
stats_details(params=params, flops=flops, activations=activations),
)