ptflops

The pytorch backend is a legacy mode that considers nn.Modules only. It is recommended for CNNs because it provides better per-layer analytics. Do not use this backend for Transformer architectures.

Supported Layers:

• 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: RNN, LSTM, GRU, MultiheadAttention, DeformConv2d, and timm vision transformers.

Usage Tips:

• If functional-level hooks conflict with custom nn.Module hooks, disable them using backend_specific_config={"count_functional": False}.
• To handle models with multiple or optional inputs, use the input_constructor argument. This is a function that takes the input spatial resolution as a tuple and returns a dictionary of named input arguments.
• Use ignore_modules to skip specific layers (e.g., ignore_modules=[torch.nn.Conv2d]).
• Use verbose=True to see modules that do not contribute to the final complexity count.

Use get_model_complexity_info to compute the theoretical multiply-add operations (MACs) and the number of parameters in a PyTorch model.

Arguments:

• net: The PyTorch model instance.
• input_size: A tuple representing the input spatial resolution (e.g., (3, 224, 224)).
• as_strings: If True, returns the results as human-readable strings.
• backend: Either 'aten' (default) or 'pytorch'.
• print_per_layer_stat: If True, prints per-layer computational costs.
• verbose: If True, prints information about uncounted operations/modules.

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))

The get_flops_aten function provides a high-level interface to estimate the MACs (Multiply-Accumulate operations) and total parameter count of a PyTorch model using the aten backend.

Arguments:

• model: The PyTorch model to analyze.
• input_res: A tuple representing the input resolution (e.g., (3, 224, 224)).
• print_per_layer_stat: Boolean, defaults to True. If True, prints layer-wise statistics to the output stream.
• input_constructor: An optional callable that takes input_res and returns a valid input batch (e.g., a tensor or a dictionary of tensors).
• ost: The output stream for printing (defaults to sys.stdout).
• verbose: Boolean, defaults to False. Enables warnings for ignored or zero-op operations.
• ignore_modules: A list of modules to ignore during the counting process.
• custom_modules_hooks: A dictionary of custom hooks for specific modules.
• output_precision: Number of decimal places for printed values.
• flops_units: The unit for FLOPs (e.g., 'GMac', 'MMac').
• param_units: The unit for parameters (e.g., 'M').
• extra_config: Additional configuration dictionary.

Returns: A tuple containing (macs_count, params_sum). If an error occurs during estimation, it returns (None, None).

The primary entry point for estimating the computational complexity of a PyTorch model. It calculates both FLOPs (Floating Point Operations) and the number of parameters.

Parameters:

• model: The PyTorch nn.Module to analyze.
• input_res: A tuple representing the input resolution (e.g., (3, 224, 224)).
• print_per_layer_stat (bool): If True, prints a detailed breakdown of FLOPs and parameters per layer.
• input_constructor (callable, optional): A function that takes input_res and returns a valid input batch (e.g., a dict or torch.Tensor). If None, a tensor of ones is used.
• ost (file-like, optional): Output stream for printing statistics (defaults to sys.stdout).
• verbose (bool): If True, prints warnings for unsupported modules treated as zero-ops.
• ignore_modules (list): A list of module types to ignore during counting.
• custom_modules_hooks (dict): A mapping of module types to custom FLOPs counting hook functions.
• output_precision (int): Decimal precision for printed statistics.
• flops_units (str): Unit for FLOPs (e.g., 'GMac').
• param_units (str): Unit for parameters (e.g., 'M').
• extra_config (dict): Configuration dictionary. Use {'count_functional': True} to enable patching of torch.nn.functional and torch.tensor operations.

Returns:

• A tuple containing (int(flops_count), params_count). Returns (None, None) if an exception occurs during estimation.