Statistics and visualization of parameters and calculations#

With some tools, we can MegEngine parameters and computation models to statistics, there are two ways to achieve：

Based on the implementation of: py:mod:~.module——
- Advantage：can be embedded and called in Python code, and statistics can be viewed at any time
- Disadvantage：can only count the information of module, but cannot count the calls of functional
Based on the implementation of: py:meth:~.trace.dump——
- Advantage：can cover all operators
- Disadvantage：needs to be performed first: py:meth:~.trace.dump operation

Statistics based on module#

Achieved: py: func: ~ .module_stats can support statistical float32 / qat / qint8 model, use very simple：

from megengine.hub import load
from megengine.utils.module_stats import module_stats

# 构建一个 net module，这里从 model hub 中获取 resnet18 模型
model = load("megengine/models", "resnet18", pretrained=True)

# 指定输入
input_data = np.random.rand(1, 3, 224, 224).astype("float32")

# Float model.
total_stats, stats_details = module_stats(
    model,
    inputs=input_data,
    cal_params=True,
    cal_flops=True,
    cal_activations=True,
    logging_to_stdout=True,
)
print("params {} flops {} acts {}".format(total_stats.param_dims, total_stats.flops, total_stats.act_dims))

You can use cal_params, cal_flops and cal_activations to control whether to calculate the parameter, flops and activations information, and use logging_to_stdout to control whether to print out the calculated details and return the total The namedtuple of statistics and detailed statistics, you can view the total amount of each statistic and the weight of each module.

Visualization and statistics based on dump graph#

Based Python Graph showing the structure analysis function implemented：

Enter the dump model path and log storage directory in mge format
The graph structure information can be saved in a format readable by TensorBoard.

Command line call#

python -m megengine.tools.network_visualize ./resnet18.mge --log_path ./log --load_input_data data.pkl --cal_flops --cal_params --cal_activations --logging_to_stdout

The description of each parameter is as follows：

./resnet18.mge (the first parameter): Required parameter, specify the model file name.
./log (the second parameter): Required parameter, specify the log storage directory.
``–load_input_data’’: Specify the path of the input data file. The content of the file should be a pickled numpy array or a dict containing a numpy array, and the key should be the input node name.
``–cal_flops’’: Specify statistics about FLOPs.
``–cal_params’’: Specify statistical Parameters information.
--cal_activations: Specify statistical activations information.
``–logging_to_stdout’’: Specify the current screen to print out all statistics information.

Call in Python#

The following code is equivalent to the command line call method above：

from megengine.tools.network_visualize import visualize

input_data = np.random.rand(1, 3, 224, 224).astype("float32")
total_stats, stats_details = visualize(
    "./resnet18.mge",
    "./log",
    input=input_data,
    cal_flops=True,
    cal_params=True,
    cal_activations=True,
    logging_to_stdout=True
)
print("params {} flops {} acts {}".format(total_stats.param_dims, total_stats.flops, total_stats.act_dims))

Visualize#

After completing the above steps, start tensorboard in the corresponding directory (in the example, ./log), and you can open the tensorboard process on this machine.：

tensorboard --logdir ./log

Note

For the installation and use of TensorBoard, please refer to TensorBoard official website.

If the startup server is a remote ssh login, the following command can be used to map the port to the local (you can use the abbreviation of the server name in sshconfig)：

ssh <user>@<host_name> -L 6006:0.0.0.0:6006 -N