torchstat: The Pytorch Model Analyzer.
Project description
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
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
$ torchstat --file example.py --model Net
[MAdd]: Dropout2d is not supported!
[Flops]: Dropout2d is not supported!
module name input shape output shape params memory(MB) MAdd Flops duration[%]
0 conv1 3 224 224 10 220 220 760.0 1.85 72,600,000.0 36,784,000.0 60.11%
1 conv2 10 110 110 20 106 106 5020.0 0.86 112,360,000.0 56,404,720.0 35.08%
2 conv2_drop 20 106 106 20 106 106 0.0 0.86 0.0 0.0 0.34%
3 fc1 56180 50 2809050.0 0.00 5,617,950.0 2,809,000.0 4.25%
4 fc2 50 10 510.0 0.00 990.0 500.0 0.22%
total 2815340.0 3.56 190,578,940.0 95,998,220.0 100.00%
==========================================================================================================
Total params: 2,815,340
----------------------------------------------------------------------------------------------------------
Total memory: 3.56MB
Total MAdd: 190.58MMAdd
Total Flops: 96.0MFlops
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
from torchvision.models as models
model = models.resnet18()
stat(model, (3, 224, 224))
Features
Note: These features work only nn.Module. Modules in torch.nn.functional are not supported yet.
- FLOPs
- Number of Parameters
- Total memory
- Madd(FMA)
- Model summary(detail, layer-wise)
- Export score table
- MemRead
- MemWrite
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.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distributions
Built Distributions
Hashes for torchstat-0.0.5-py3-none-any.whl
Algorithm | Hash digest | |
---|---|---|
SHA256 | c33e442c78e6c41016d6bda0c1e00c71177746eb808d5d861a8f8b417c8a376a |
|
MD5 | 7e959fd22e1c27a99e95419f18f5fa2a |
|
BLAKE2b-256 | 23923de1c565506aa63789f77cfcfc122626279ae74eac57fe053f1b4390fe24 |