热搜：编程 Python Microsoft 编程语言 C++

pytorch 计算Parameter和FLOP的操作

发布时间：2022-12-27 01:50:03 来源：好代码

.

挂林风景异，秋似洛阳春。倘若，生命是一场旅行，爱，一定在路上。而往往，走得最急的，都是最美的风景。世界上有不绝的风景，我有不老的心。有时三点两点雨，到处十枝五枝花。

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

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, 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)