pytorch 计算Parameter和FLOP的操作

(编辑:jimmy 日期: 2025/2/25 浏览:2)

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

1 THOP

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

pip install thop

注意安装thop时可能出现如下错误:

pytorch 计算Parameter和FLOP的操作

解决方法:

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, 224, 224),) 
  print("%.2fM" % (flop/1e6), "%.2fM" % (para/1e6))

更多细节,可参考thop GitHub链接: https://github.com/Lyken17/pytorch-OpCounter

2 计算参数

pytorch本身带有计算参数的方法

  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)
  total = sum([param.nelement() for param in model.parameters()])
  print("Number of parameter: %.2fM" % (total / 1e6))

补充:pytorch: 计算网络模型的计算量(FLOPs)和参数量(Params)

计算量:

FLOPs,FLOP时指浮点运算次数,s是指秒,即每秒浮点运算次数的意思,考量一个网络模型的计算量的标准。

参数量:

Params,是指网络模型中需要训练的参数总数。

第一步:安装模块(thop)

pip install thop

第二步:计算

import torch
from thop import profile
net = Model() # 定义好的网络模型
input = torch.randn(1, 3, 112, 112)
flops, params = profile(net, (inputs,))
print('flops: ', flops, 'params: ', params)

注意:

输入input的第一维度是批量(batch size),批量的大小不回影响参数量, 计算量是batch_size=1的倍数

profile(net, (inputs,))的 (inputs,)中必须加上逗号,否者会报错

以上为个人经验,希望能给大家一个参考,也希望大家多多支持。如有错误或未考虑完全的地方,望不吝赐教。

一句话新闻
一文看懂荣耀MagicBook Pro 16
荣耀猎人回归!七大亮点看懂不只是轻薄本,更是游戏本的MagicBook Pro 16.
人们对于笔记本电脑有一个固有印象:要么轻薄但性能一般,要么性能强劲但笨重臃肿。然而,今年荣耀新推出的MagicBook Pro 16刷新了人们的认知——发布会上,荣耀宣布猎人游戏本正式回归,称其继承了荣耀 HUNTER 基因,并自信地为其打出“轻薄本,更是游戏本”的口号。
众所周知,寻求轻薄本的用户普遍更看重便携性、外观造型、静谧性和打字办公等用机体验,而寻求游戏本的用户则普遍更看重硬件配置、性能释放等硬核指标。把两个看似难以相干的产品融合到一起,我们不禁对它产生了强烈的好奇:作为代表荣耀猎人游戏本的跨界新物种,它究竟做了哪些平衡以兼顾不同人群的各类需求呢?