深度学习中,模型训练完后,查看模型的参数量和浮点计算量,在此记录下:
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