SoFunction
Updated on 2025-03-04

Basic use cases and parameter description of Thop library in Python

Introduction to Thop

THOP (Torch-OpCounter) is a tool used to calculate the operands and calculation amount of PyTorch model. Through this library, developers can better understand and evaluate the complexity of the model, which is very helpful for model optimization and performance tuning.

Specifically, using THOP, the defined PyTorch model can be analyzed and the number of various operations contained in the model and the amount of calculation of the model can be counted. This is very important for understanding the operating efficiency of the model, identifying possible performance bottlenecks, and optimizing the model structure.

The basic steps of using the THOP library in Python include: first install the THOP library, which can usually be installed through the pip command; then import the THOP library in Python scripts; finally define the PyTorch model that needs to be analyzed, and use THOP to statistics on the operands and calculation amount of the model.

It should be noted that the specific usage methods and supported functions of the THOP library may change with the update of the version, so it is recommended to consult the latest official documentation for the most accurate information.

parameter

  • model(PyTorch model): A PyTorch model that requires calculation of operands and calculation amounts.
  • input_size(Input Size): The size of the input tensor of the model. This is usually a tuple that specifies the shape of the input tensor.
  • custom_ops(Custom Action): A dictionary that specifies the amount of calculation for a custom operation. This can be used to override the default operation calculation method in the THOP library.
  • ignore_ops(Ignore operations): A list that specifies the types of operations that need to be ignored during the calculation process. This is useful for excluding some unimportant operations or computational volumes of specific layers.

Basic use cases

  • Install the THOP library: First, you need to install the THOP library. Usually, you can use the pip command to install:
pip install thop
  • Import the necessary libraries: In your Python script, import the necessary libraries and modules:
import torch  
import thop
  • Defining the PyTorch model: Define a PyTorch model that will be used to calculate operands and calculation amounts.
  • Calculate operands and calculation amount: Using the THOP libraryprofileFunctions to calculate the operands and calculation amount of the model. You need to provide the model and input size as parameters. For example:
model = ...  # Define your PyTorch modelinput_size = (1, 3, 224, 224)  # Enter the size, here assume that it is a batch of 3-channel 224x224 images  
flops, params = (model, inputs=((input_size),))  
print(f"FLOPs: {flops / 1e9} G")  # Print calculations (in billion floating point operations)print(f"Params: {params / 1e6} M")  # Print parameter quantity (in million units)

In this example,The function returns two values:flopsRepresents the calculation amount of the model (number of floating-point operations),paramsRepresents the number of parameters of the model. These two values ​​are often used to evaluate the complexity and performance of the model. Note that when calculating FLOPs, we usually divide them by1e9to convert it into units of billion floating point operations (GFLOPs), and likewise, the parameter quantity is usually divided by1e6to convert it into units of millions (MParams).

Attachment: thop installation problems

Friends who have installed third-party libraries in python know that its installation command is:pip install XXX. However, there is a big pit for Thop installation.

1) Directly use common commands

pip install thop

It will install the version 0.0.31.post2005241907 by default. Since the pytorch environment I'm using is1.0.0, not supported:(necessary required for multi-GPU execution), and the THop under this version【THOP:0.031] This bug is not fixed for pytorch:1.0.0, resulting in a running error.

2) Use git library installation method

pip install --upgrade git+/Lyken17/

The second way to install will be the result of the latest Thop library! ! ! Super convenient! ! !

Summarize

This is the article about the basic use cases and parameter descriptions of Thop libraries in Python. For more related content on using Thop libraries in Python, please search for my previous articles or continue browsing the related articles below. I hope everyone will support me in the future!