calculate-flops.pytorch

calflops: A Tool for Calculating FLOPs and Parameters for Neural Networks

Overview

calflops is a specialized tool designed to calculate the theoretical number of FLOPs (Floating Point Operations), MACs (Multiply-Add Operations), and parameters required by a variety of neural network models. It supports computation for numerous architectures including Linear models, CNNs (Convolutional Neural Networks), RNNs (Recurrent Neural Networks), GCNs (Graph Convolutional Networks), and Transformers such as Bert and LLaMA. Additionally, it can accommodate custom models developed with PyTorch using torch.nn.functional.*.

Key Features

• Comprehensive Compatibility: It covers a wide range of neural network types, including large language models and any custom models built with PyTorch.
• Detailed Output: Users can print not only the total FLOPs, MACs, and parameters but also the breakdown and proportion of these metrics for each submodule of a model, providing better insight into which parts of the model are performance bottlenecks.
• Integration with Hugging Face: Calflops is now available as a tool on Hugging Face Space, making it easier to compute FLOPs directly on models hosted there without downloading them.

Usage

Calflops can be installed via PyPI with the following command:

pip install --upgrade calflops

Example Usages:

1. CNN Model: For calculating FLOPs in a CNN model like AlexNet:

   from calflops import calculate_flops
   from torchvision import models
   
   model = models.alexnet()
   flops, macs, params = calculate_flops(model=model, input_shape=(1, 3, 224, 224))
   print(f"AlexNet FLOPs: {flops}, MACs: {macs}, Params: {params}")
   

2. Hugging Face Models: For models on Hugging Face, calflops can compute FLOPs via model names without downloading full weights:

   from calflops import calculate_flops_hf
   
   flops, macs, params = calculate_flops_hf(model_name="meta-llama/Llama-2-7b", input_shape=(1, 128))
   print(f"Model FLOPs: {flops}, MACs: {macs}, Params: {params}")
   

3. Transformer and LLMs: When dealing with Transformer models or large language models, users can specify tokenizers to automatically handle input preparation:

   from transformers import AutoModel, AutoTokenizer
   from calflops import calculate_flops
   
   model = AutoModel.from_pretrained("bert-base-uncased")
   tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
   flops, macs, params = calculate_flops(model=model, transformer_tokenizer=tokenizer, input_shape=(1, 128))
   print(f"Bert FLOPs: {flops}, MACs: {macs}, Params: {params}")
   

Advanced Features

• Customizable Output: Users can refine the output format, control the precision of results, and select units like TFLOPs or GFLOPs.
• Handling Multiple Inputs: The tool supports models with multiple input data streams by using args or kwargs.
• Backward Pass: Flops calculation can optionally include both forward and backward passes required for model training.
• Selective Module Evaluation: Users can ignore specific modules during the computational evaluation to focus on particular parts of a model.

Inspiration and Development

Calflops draws inspiration from existing libraries such as ptflops, DeepSpeed, and Hugging Face's accelerate. Its development involves improvements that enhance FLOPs calculation accuracy and efficiency.

In summary, calflops is a versatile and easy-to-use tool for assessing the computational demands of neural network models. It serves as a valuable resource for developers and researchers aiming to optimize model architecture without delving into manual calculations.