pytorch 計算Parameter和FLOP的操作

Posted on by

深度學習中，模型訓練完後，查看模型的參數量和浮點計算量，在此記錄下：

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)

註意：

輸入input的第一維度是批量(batch size)，批量的大小不回影響參數量， 計算量是batch_size=1的倍數

profile(net, (inputs,))的 (inputs,)中必須加上逗號，否者會報錯

以上為個人經驗，希望能給大傢一個參考，也希望大傢多多支持WalkonNet。如有錯誤或未考慮完全的地方，望不吝賜教。

推薦閱讀：