torchstat

torchstat

This is a lightweight neural network analyzer based on PyTorch. It is designed to make building your networks quick and easy, with the ability to debug them. Note: This repository is currently under development. Therefore, some APIs might be changed.

This tools can show

• Total number of network parameters
• Theoretical amount of floating point arithmetics (FLOPs)
• Theoretical amount of multiply-adds (MAdd)
• Memory usage

Installing

There're two ways to install torchstat into your environment.

• Install it via pip.

$ pip install torchstat

• Install and update using setup.py after cloning this repository.

$ python3 setup.py install

A Simple Example

If you want to run the torchstat asap, you can call it as a CLI tool if your network exists in a script. Otherwise you need to import torchstat as a module.

CLI tool

…

If you're not sure how to use a specific command, run the command with the -h or –help switches. You'll see usage information and a list of options you can use with the command.

Module

from torchstat import stat
import torchvision.models as models

model = models.resnet18()
stat(model, (3, 224, 224))

Features & TODO

Note: These features work only nn.Module. Modules in torch.nn.functional are not supported yet.

• FLOPs
• Number of Parameters
• Total memory
• Madd(FMA)
• MemRead
• MemWrite
• Model summary(detail, layer-wise)
• Export score table
• Arbitrary input shape

For the supported layers, check out the details.

Requirements

• Python 3.6+
• Pytorch 0.4.0+
• Pandas 0.23.4+
• NumPy 1.14.3+

References

Thanks to @sovrasov for the initial version of flops computation, @ceykmc for the backbone of scripts.

• flops-counter.pytorch
• pytorch_model_summary
• chainer_computational_cost
• convnet-burden.