OneFlow-OpCounter | 简体中文

modified from https://github.com/sovrasov/flops-counter.pytorch

install

pip install flowflops

usage

importoneflowasflowfromflowflopsimportget_model_complexity_infofromflowflops.utilsimportflops_to_string, params_to_stringmodel= ... # your own model, nn.Moduledsize= (1, 3, 224, 224) # B, C, H, Wtotal_flops, total_params=get_model_complexity_info(
model, dsize,
as_strings=False,
print_per_layer_stat=False,
mode="eager"# eager or graph
)
print(flops_to_string(total_flops), params_to_string(total_params))

why graph?

classBasicBlock(nn.Module):
expansion: int=1def__init__(
self,
inplanes: int,
planes: int,
stride: int=1,
downsample: Optional[nn.Module] =None,
groups: int=1,
base_width: int=64,
dilation: int=1,
norm_layer: Optional[Callable[..., nn.Module]] =None,
) ->None:
super(BasicBlock, self).__init__()
ifnorm_layerisNone:
norm_layer=nn.BatchNorm2difgroups!=1orbase_width!=64:
raiseValueError("BasicBlock only supports groups=1 and base_width=64")
ifdilation>1:
raiseNotImplementedError("Dilation > 1 not supported in BasicBlock")
# Both self.conv1 and self.downsample layers downsample the input when stride != 1self.conv1=conv3x3(inplanes, planes, stride)
self.bn1=norm_layer(planes)
self.relu=nn.ReLU()
self.conv2=conv3x3(planes, planes)
self.bn2=norm_layer(planes)
self.downsample=downsampleself.stride=stridedefforward(self, x: Tensor) ->Tensor:
identity=xout=self.conv1(x)
out=self.bn1(out)
out=self.relu(out)
out=self.conv2(out)
out=self.bn2(out)
ifself.downsampleisnotNone:
identity=self.downsample(x)
out+=identity# !!!NOTE!!!: this will make add-op flops that cannot be hooked in eager modeout=self.relu(out)
returnout

sample

python benchmark/evaluate_famous_models.py

====== eager ======
+--------------------+----------+-------------+
| Model | Params | FLOPs |
+--------------------+----------+-------------+
| alexnet | 61.1 M | 718.16 MMac |
| vgg11 | 132.86 M | 7.63 GMac |
| vgg11_bn | 132.87 M | 7.64 GMac |
| squeezenet1_0 | 1.25 M | 830.05 MMac |
| squeezenet1_1 | 1.24 M | 355.86 MMac |
| resnet18 | 11.69 M | 1.82 GMac |
| resnet50 | 25.56 M | 4.12 GMac |
| resnext50_32x4d | 25.03 M | 4.27 GMac |
| shufflenet_v2_x0_5 | 1.37 M | 43.65 MMac |
| regnet_x_16gf | 54.28 M | 16.01 GMac |
| efficientnet_b0 | 5.29 M | 401.67 MMac |
| densenet121 | 7.98 M | 2.88 GMac |
+--------------------+----------+-------------+
====== graph ======
+--------------------+----------+-------------+
| Model | Params | FLOPs |
+--------------------+----------+-------------+
| alexnet | 61.1 M | 718.16 MMac |
| vgg11 | 132.86 M | 7.63 GMac |
| vgg11_bn | 132.87 M | 7.64 GMac |
| squeezenet1_0 | 1.25 M | 830.05 MMac |
| squeezenet1_1 | 1.24 M | 355.86 MMac |
| resnet18 | 11.69 M | 1.82 GMac |
| resnet50 | 25.56 M | 4.13 GMac |
| resnext50_32x4d | 25.03 M | 4.28 GMac |
| shufflenet_v2_x0_5 | 1.37 M | 43.7 MMac |
| regnet_x_16gf | 54.28 M | 16.02 GMac |
| efficientnet_b0 | 5.29 M | 410.35 MMac |
| densenet121 | 7.98 M | 2.88 GMac |
+--------------------+----------+-------------+

support

Eager

the outputs will be the same as the ptflops

supported layers:

# convolutionsnn.Conv1dnn.Conv2dnn.Conv3d# activationsnn.ReLUnn.PReLUnn.ELUnn.LeakyReLUnn.ReLU6# poolingsnn.MaxPool1dnn.AvgPool1dnn.AvgPool2dnn.MaxPool2dnn.MaxPool3dnn.AvgPool3d# nn.AdaptiveMaxPool1dnn.AdaptiveAvgPool1d# nn.AdaptiveMaxPool2dnn.AdaptiveAvgPool2d# nn.AdaptiveMaxPool3dnn.AdaptiveAvgPool3d# BNsnn.BatchNorm1dnn.BatchNorm2dnn.BatchNorm3d# INsnn.InstanceNorm1dnn.InstanceNorm2dnn.InstanceNorm3d# FCnn.Linear# Upscalenn.Upsample# Deconvolutionnn.ConvTranspose1dnn.ConvTranspose2dnn.ConvTranspose3d# RNNnn.RNNnn.GRUnn.LSTMnn.RNNCellnn.LSTMCellnn.GRUCell

Graph

supported ops:

# conv"conv1d""conv2d""conv3d"# pool"max_pool_1d""max_pool_2d""max_pool_3d""avg_pool_1d""avg_pool_2d""avg_pool_3d""adaptive_max_pool1d""adaptive_max_pool2d""adaptive_max_pool3d""adaptive_avg_pool1d""adaptive_avg_pool2d""adaptive_avg_pool3d"# activate"relu""leaky_relu""prelu""hardtanh""elu""silu""sigmoid""sigmoid_v2"# add"bias_add""add_n"# matmul"matmul""broadcast_matmul"# norm"normalization"# scalar"scalar_mul""scalar_add""scalar_sub""scalar_div"# stats"var"# math"sqrt""reduce_sum"# broadcast"broadcast_mul""broadcast_add""broadcast_sub""broadcast_div"# empty"reshape""ones_like""zero_like""flatten""concat""transpose""slice"