THOP：浮点运算量（ FLOPs）和参数数量（Params）的计算工具

一、简介

thop是一种常用的估算PyTorch模型FLOPs和参数数量的工具，它通过分析模型的计算图，自动计算模型进行一次前向推理时每一层的操作次数和参数数量，并且支持多种类型的神经网络结构。FLOPs是衡量模型计算复杂度的一个重要指标，FLOPs越高，通常模型的计算量也越大。而参数数量是模型中所有可训练参数的总和，计算参数数量有助于评估模型的存储需求。

二、安装及使用

1. 安装thop。

   pip install thop
   

2. thop主要通过profile函数来分析模型的FLOPs和参数数量。

   from thop import profile
   ------------------------------------------------
   # 假设你有一个 PyTorch 模型 `model` 和输入张量 `x`
   flops, params = profile(model, inputs=(x,))
   

3. 增加更多解释性的输出文本，thop输出的FLOPs单位通常是浮动的，你可以选择以GFlops（十亿次浮点运算）或TFlops（万亿次浮点运算）为单位来显示。

   flops, params = profile(model, (x,))
   
   # 输出 FLOPs 和参数数量，增加描述文字
   print(f"Model Complexity:")
   print(f"  - FLOPs: {flops / 1e9:.2f} GFlops (billion floating point operations)")
   print(f"  - Parameters: {params / 1e6:.2f} M (million parameters)")
   

4. 加入计算时长的计算。

   import time
   
   start_time = time.time()
   
   model(x)
   
   end_time = time.time()
   
   # 计算时间差
   elapsed_time = end_time - start_time
   
   # 计算复杂度
   flops, params = profile(model, (x,))
   print(f"Model Complexity and compute time:")
   print(f"  - FLOPs:  {flops / 1e9:.2f} GFlops   (billion floating point operations)")
   print(f"  - Params: {params / 1e6:.2f} M       (million parameters)")
   print(f"  - Time:   {elapsed_time:.2f} S       (seconds of a forward pass)")
   

5. 另外，可以使用torchsummary来显示模型的详细信息。

   pip install torchsummary
   ------------------------------------------------
   from torchsummary import summary
   
   summary(model, (3, 224, 224))  # 输入张量的形状
   

   

小结

thop是一个非常轻量且高效的工具，适合快速计算PyTorch模型的FLOPs和参数数量。通过它，可以快速了解模型的计算复杂度和存储需求，帮助在开发过程中做出合理的优化和选择。在实际应用中，除了thop，还可以结合其他工具（如torchsummary）来进一步分析模型的性能和结构。