🚀 THOP: PyTorch-OpCounter

Welcome to the THOP repository, your comprehensive solution for profiling PyTorch models by computing the number of Multiply-Accumulate Operations (MACs) and parameters. Developed by Ultralytics, this tool is essential for deep learning practitioners aiming to evaluate model efficiency and performance, crucial aspects discussed in our model training tips guide.

📄 Description

THOP offers an intuitive API designed to profile PyTorch models by calculating the total number of MACs and parameters. This functionality is vital for assessing the computational efficiency and memory footprint of deep learning models, helping developers optimize performance for deployment, especially on edge devices. Understanding these metrics is key to selecting the right model architecture, a topic explored in our model comparison pages.

📦 Installation

Get started with THOP quickly by installing it via pip:

pip install ultralytics-thop

Alternatively, for the latest features and updates, install directly from the GitHub repository:

pip install --upgrade git+https://github.com/ultralytics/thop.git

This ensures you have the most recent version, incorporating the latest improvements and bug fixes.

🛠️ How to Use

Basic Usage

Profiling a standard PyTorch model like ResNet50 is straightforward. Import the necessary libraries, load your model and a sample input tensor, then use the profile function:

import torch
from torchvision.models import resnet50  # Example model

from thop import profile  # Import the profile function from THOP

# Load a pre-trained model (e.g., ResNet50)
model = resnet50()

# Create a dummy input tensor matching the model's expected input shape
dummy_input = torch.randn(1, 3, 224, 224)

# Profile the model
macs, params = profile(model, inputs=(dummy_input,))

print(f"MACs: {macs}, Parameters: {params}")
# Expected output: MACs: 4139975680.0, Parameters: 25557032.0

Define Custom Rules for Third-Party Modules

If your model includes custom or third-party modules not natively supported by THOP, you can define custom profiling rules using the custom_ops argument. This allows for accurate profiling even with complex or non-standard architectures, which is useful when working with models like those found in the Ultralytics models section.

import torch
import torch.nn as nn

from thop import profile


# Define your custom module
class YourCustomModule(nn.Module):
    def __init__(self):
        super().__init__()
        # Define layers, e.g., a convolution
        self.conv = nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)

    def forward(self, x):
        return self.conv(x)


# Define a custom counting function for your module
# This function should calculate and return the MACs for the module's operations
def count_your_custom_module(module, x, y):
    # Example: Calculate MACs for the conv layer
    # Note: This is a simplified example. Real calculations depend on the module's specifics.
    # MACs = output_height * output_width * kernel_height * kernel_width * in_channels * out_channels
    # For simplicity, we'll just assign a placeholder value or use a helper if available
    # In a real scenario, you'd implement the precise MAC calculation here.
    # For nn.Conv2d, THOP usually handles it, but this demonstrates the concept.
    macs = 0  # Placeholder: Implement actual MAC calculation based on module logic
    # You might need access to module properties like kernel_size, stride, padding, channels etc.
    # Example for a Conv2d layer (simplified):
    if isinstance(module, nn.Conv2d):
        _, _, H, W = y.shape  # Output shape
        k_h, k_w = module.kernel_size
        in_c = module.in_channels
        out_c = module.out_channels
        groups = module.groups
        macs = (k_h * k_w * in_c * out_c * H * W) / groups
    module.total_ops += torch.DoubleTensor([macs])  # Accumulate MACs


# Instantiate a model containing your custom module
model = YourCustomModule()  # Or a larger model incorporating this module

# Create a dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# Profile the model, providing the custom operation mapping
macs, params = profile(model, inputs=(dummy_input,), custom_ops={YourCustomModule: count_your_custom_module})

print(f"Custom MACs: {macs}, Parameters: {params}")
# Expected output: Custom MACs: 87457792.0, Parameters: 1792.0

Improve Output Readability

For clearer and more interpretable results, use the thop.clever_format function. This formats the raw MACs and parameter counts into human-readable strings (e.g., GigaMACs, MegaParams). This formatting helps in quickly understanding the scale of computational resources required, similar to the metrics provided in our Ultralytics YOLOv8 documentation.

import torch
from torchvision.models import resnet50

from thop import clever_format, profile

model = resnet50()
dummy_input = torch.randn(1, 3, 224, 224)
macs, params = profile(model, inputs=(dummy_input,))

# Format the numbers into a readable format (e.g., 4.14 GMac, 25.56 MParams)
macs_readable, params_readable = clever_format([macs, params], "%.3f")

print(f"Formatted MACs: {macs_readable}, Formatted Parameters: {params_readable}")
# Expected output: Formatted MACs: 4.140G, Formatted Parameters: 25.557M

📊 Results of Recent Models

The table below showcases the parameters and MACs for several popular computer vision models, profiled using THOP. These benchmarks provide a comparative overview of model complexity and computational cost. You can reproduce these results by running the script located at benchmark/evaluate_famous_models.py in this repository. Comparing these metrics is essential for tasks like selecting models for object detection or image classification. For more comparisons, see our model comparison section.

🙌 Contribute

We actively welcome and encourage community contributions to make THOP even better! Whether it's adding support for new PyTorch layers, improving existing calculations, enhancing documentation, or fixing bugs, your input is valuable. Please see our Contributing Guide for detailed instructions on how to participate. Together, we can ensure THOP remains a state-of-the-art tool for the machine learning community. Don't hesitate to share your feedback and suggestions!

📜 License

THOP is distributed under the AGPL-3.0 License. This license promotes open collaboration and sharing of improvements. For complete details, please refer to the LICENSE file included in the repository. Understanding the license is important before integrating THOP into your projects, especially for commercial applications which may require an Enterprise License.

📧 Contact

Encountered a bug or have a feature request? Please submit an issue through our GitHub Issues page. For general discussions, questions, and community support, join the vibrant Ultralytics community on our Discord server. We look forward to hearing from you and collaborating!