Flops counter for convolutional networks in pytorch framework

This script is designed to compute the theoretical amount of multiply-add operations in convolutional neural networks. It also can compute the number of parameters and print per-layer computational cost of a given network.

Supported layers:

• Conv1d/2d/3d (including grouping)
• ConvTranspose1d/2d/3d (including grouping)
• BatchNorm1d/2d/3d
• Activations (ReLU, PReLU, ELU, ReLU6, LeakyReLU)
• Linear
• Upsample
• Poolings (AvgPool1d/2d/3d, MaxPool1d/2d/3d and adaptive ones)

Experimental support:

• RNN, LSTM, GRU (NLH layout is assumed).
• RNNCell, LSTMCell, GRUCell.

Requirements: Pytorch >= 0.4.1, torchvision >= 0.2.1

Thanks to @warmspringwinds for the initial version of script.

Usage tips

• This script doesn't take into account torch.nn.functional.* operations. For an instance, if one have a semantic segmentation model and use torch.nn.functional.interpolate to upscale features, these operations won't contribute to overall amount of flops. To avoid that one can use torch.nn.Upsample instead of torch.nn.functional.interpolate.
• ptflops launches a given model on a random tensor and estimates amount of computations during inference. Complicated models can have several inputs, some of them could be optional. To construct non-trivial input one can use the input_constructor argument of the get_model_complexity_info. input_constructor is a function that takes the input spatial resolution as a tuple and returns a dict with named input arguments of the model. Next this dict would be passed to the model as keyworded arguments.
• verbose parameter allows to get information about modules that don't contribute to the final numbers.
• ignore_modules option forces ptflops to ignore the listed modules. This can be useful for research purposes. For an instance, one can drop all convolutuions from the counting process specifying ignore_modules=[torch.nn.Conv2d].

Install the latest version

From PyPI:

pip install ptflops

From this repository:

pip install --upgrade git+https://github.com/sovrasov/flops-counter.pytorch.git

Example

import torchvision.models as models
import torch
from ptflops import get_model_complexity_info

with torch.cuda.device(0):
  net = models.densenet161()
  macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True,
                                           print_per_layer_stat=True, verbose=True)
  print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
  print('{:<30}  {:<8}'.format('Number of parameters: ', params))

Benchmark

torchvision

• Top-1 error - ImageNet single-crop top-1 error (224x224)
• Top-5 error - ImageNet single-crop top-5 error (224x224)

Cadene/pretrained-models.pytorch

• [email protected] - ImageNet single-crop top-1 accuracy on validation images of the same size used during the training process.
• [email protected] - ImageNet single-crop top-5 accuracy on validation images of the same size used during the training process.