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

ptflops has two backends, pytorch and aten. pytorch backend is a legacy one, it considers nn.Modules only. However, it's still useful, since it provides a better par-layer analytics for CNNs. In all other cases it's recommended to use aten backend, which considers aten operations, and therefore it covers more model architectures (including transformers).

• aten.mm, aten.matmul, aten.addmm, aten.bmm
• aten.convolution

• Use verbose=True to see the operations which were not considered during complexity computation.
• This backend prints per-module statistics only for modules directly nested into the root nn.Module. Deeper modules at the second level of nesting are not shown in the per-layer statistics.

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

Experimental support:

• RNN, LSTM, GRU (NLH layout is assumed)
• RNNCell, LSTMCell, GRUCell
• torch.nn.MultiheadAttention
• torchvision.ops.DeformConv2d
• visual transformers from timm

• This backend doesn't take into account some of the torch.nn.functional.* and tensor.* operations. Therefore unsupported operations are not contributing to the final complexity estimation. See ptflops/pytorch_ops.py:FUNCTIONAL_MAPPING,TENSOR_OPS_MAPPING to check supported ops.
• 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 a keyword 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 instance, one can drop all convolutions from the counting process specifying ignore_modules=[torch.nn.Conv2d].

Requirements: Pytorch >= 1.1, torchvision >= 0.3

Thanks to @warmspringwinds and Horace He for the initial version of the script.

From PyPI:

pip install ptflops

From this repository:

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

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, backend='pytorch'
                                           print_per_layer_stat=True, verbose=True)
  print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
  print('{:<30}  {:<8}'.format('Number of parameters: ', params))

  macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, backend='aten'
                                           print_per_layer_stat=True, verbose=True)
  print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
  print('{:<30}  {:<8}'.format('Number of parameters: ', params))

If ptflops was useful for your paper or tech report, please cite me:

@online{ptflops,
  author = {Vladislav Sovrasov},
  title = {ptflops: a flops counting tool for neural networks in pytorch framework},
  year = 2018-2024,
  url = {https://github.com/sovrasov/flops-counter.pytorch},
}

Model | Input Resolution | Params(M) | MACs(G)