pytorch计算Parameter和FLOP

最新推荐文章于 2025-06-30 11:27:53 发布

原创

最新推荐文章于 2025-06-30 11:27:53 发布 · 7.5k 阅读

20 ·

CC 4.0 BY-SA版权

文章标签：

#FLOP #pytorch

本文介绍了在PyTorch中如何利用thop库计算深度学习模型的参数数量和FLOPs，包括安装thop时可能遇到的问题及解决方法，以及如何对自定义函数进行计算。同时，提到了PyTorch自带的计算参数方法。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

深度学习中，模型训练完后，查看模型的参数量和浮点计算量，在此记录下：

1 THOP

在pytorch中有现成的包thop用于计算参数数量和FLOP，首先安装thop：

pip install thop

注意安装thop时可能出现如下错误：
在这里插入图片描述
解决方法：

pip install --upgrade git+https://github.com/Lyken17/pytorch-OpCounter.git  # 下载源码安装

使用方法如下：

from torchvision.models import resnet50  # 引入ResNet50模型
from thop import profile

model = resnet50()
flops, params = profile(model, input_size=(1, 3, 224,224))  #  profile（模型，输入数据）

对于自己构建的函数也一样，例如shuffleNetV2

    from thop import profile
    from utils.ShuffleNetV2 import shufflenetv2  # 导入shufflenet2 模块
    import torch 
    
    model_shuffle = shufflenetv2(width_mult=0.5)
    model = torch.nn.DataParallel(model_shuffle)   # 调用shufflenet2 模型，该模型为自己定义的
    flop, para = profile(model, input_size=(1, 3, 2