hey-yahei / OpSummary.MXNet

MXOP: MXNet-OpSummary

It only works for gluon yet.

Reference: THOP: PyTorch-OpCounter

Installation

• PyPi

  pip install --index-url https://pypi.org/simple/ mxop

• Github (latest)

  pip install --upgrade git+https://github.com/hey-yahei/OpSummary.MXNet.git

Usage

Gluon

• Count OPs for model (default)

  from mxnet import nd
  from mxop.gluon import count_ops
  inputs = (nd.ones(shape=(1, 3, 224, 224)), )
  op_counter = count_ops(net, inputs, [, exclude])   # `net` is the gluon model you want to count OPs 

• Count OPs for every block

  from mxop.gluon import count_ops
  op_counters = count_ops(net, inputs, per_block=True)
  # Print muls
  mul_counter = {k: v['muls'] for k, v in op_counters.items()}
  total = sum(mul_counter.values())
  for k, v in mul_counter.items():
      print(f"{k.name}\t{v/total*100:.2f}%")  

• Count parameters

  from mxop.gluon import count_params
  params_counter = count_params(net [, exclude])   # `net` is the gluon model you want to count parameters
                                                                  # `exclude` is the list of blocks to be excluded 

• Print summary

  from mxop.gluon import op_summary
  op_summary(net [, inputs] [, exclude])   # `net` is the gluon model you want to count
                                           # `inputs` is the inputs to feed net 
                                           # `exclude` is the list of blocks to be excluded

Custom OPs

Defining count function for your own blocks is supported.

1. Define a function hook(m, x, y) --> None in which you should set values of OPs for dict m.ops, for example,

   def count_bn(m, x, y):
       x = x[0]
       n_elem = x.size
       m.ops["adds"] = 2*n_elem
       m.ops["muls"] = n_elem
       m.ops["divs"] = n_elem
       m.ops["exps"] = 0

2. Use it as a parameter when you call function count_op or op_summary, for example,

   from mxnet.gluon import nn
   from mxop.gluon import count_ops, op_summary
   custom_ops = {nn.BatchNorm: count_bn}
   counter = count_ops(net, inputs, custom_ops=custom_ops)
   # op_summaary(net, inputs, custom_ops=custom_ops)

You can also use your own count function for blocks that I have setted count function for, because functions list in custom_ops is given higher priority.

Test

Run tests/test_gluon_utils.py to count OPs and parameters for all models in model zoo of MXNet.

Result:

To compare classification models used as backbone--
*Params col shows the number of parameters for models without last several layers.
*Muls col shows the number of multiplications for models without last several layers.

The data above may help with design when you choose such CNNs as backbone.

TODO

• Count OPs and parameters for each layer.
• Support Symbol model for MXNet.
• Support quantized models.

More details refer to 《模型参数与计算量 | Hey~YaHei!》