计算CNN 模型算力(flops)和参数(parameters)数量计算

安装基本库:
​​​​​​​pip install ptflops
pip install --upgrade git+https://github.com/sovrasov/flops-counter.pytorch.git
详细可参考: https://github.com/sovrasov/flops-counter.pytorch

代码如下:

import torchvision.models as models
import torch
from ptflops import get_model_complexity_info

with torch.cuda.device(0):
  net = models.vgg16()
  macs, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True,
                                           print_per_layer_stat=True, verbose=True)
  print('{:<30}  {:<8}'.format('Computational complexity: ', macs))
  print('{:<30}  {:<8}'.format('Number of parameters: ', params))
root@2c5bb9a7b477:/home/Semisup_cloudSemSeg_master# python Print_CNN_FLOPS_parameters.py
Warning: module Dropout is treated as a zero-op.
Warning: module VGG is treated as a zero-op.
VGG(
  138.36 M, 100.000% Params, 15.5 GMac, 100.000% MACs,
  (features): Sequential(
    14.71 M, 10.635% Params, 15.38 GMac, 99.202% MACs,
    (0): Conv2d(1.79 k, 0.001% Params, 89.92 MMac, 0.580% MACs, 3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (1): ReLU(0, 0.000% Params, 3.21 MMac, 0.021% MACs, inplace=True)
    (2): Conv2d(36.93 k, 0.027% Params, 1.85 GMac, 11.951% MACs, 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (3): ReLU(0, 0.000% Params, 3.21 MMac, 0.021% MACs, inplace=True)
    (4): MaxPool2d(0, 0.000% Params, 3.21 MMac, 0.021% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (5): Conv2d(73.86 k, 0.053% Params, 926.45 MMac, 5.976% MACs, 64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (6): ReLU(0, 0.000% Params, 1.61 MMac, 0.010% MACs, inplace=True)
    (7): Conv2d(147.58 k, 0.107% Params, 1.85 GMac, 11.941% MACs, 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (8): ReLU(0, 0.000% Params, 1.61 MMac, 0.010% MACs, inplace=True)
    (9): MaxPool2d(0, 0.000% Params, 1.61 MMac, 0.010% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (10): Conv2d(295.17 k, 0.213% Params, 925.65 MMac, 5.971% MACs, 128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (11): ReLU(0, 0.000% Params, 802.82 KMac, 0.005% MACs, inplace=True)
    (12): Conv2d(590.08 k, 0.426% Params, 1.85 GMac, 11.936% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (13): ReLU(0, 0.000% Params, 802.82 KMac, 0.005% MACs, inplace=True)
    (14): Conv2d(590.08 k, 0.426% Params, 1.85 GMac, 11.936% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (15): ReLU(0, 0.000% Params, 802.82 KMac, 0.005% MACs, inplace=True)
    (16): MaxPool2d(0, 0.000% Params, 802.82 KMac, 0.005% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (17): Conv2d(1.18 M, 0.853% Params, 925.25 MMac, 5.968% MACs, 256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (18): ReLU(0, 0.000% Params, 401.41 KMac, 0.003% MACs, inplace=True)
    (19): Conv2d(2.36 M, 1.706% Params, 1.85 GMac, 11.933% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (20): ReLU(0, 0.000% Params, 401.41 KMac, 0.003% MACs, inplace=True)
    (21): Conv2d(2.36 M, 1.706% Params, 1.85 GMac, 11.933% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (22): ReLU(0, 0.000% Params, 401.41 KMac, 0.003% MACs, inplace=True)
    (23): MaxPool2d(0, 0.000% Params, 401.41 KMac, 0.003% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
    (24): Conv2d(2.36 M, 1.706% Params, 462.52 MMac, 2.983% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (25): ReLU(0, 0.000% Params, 100.35 KMac, 0.001% MACs, inplace=True)
    (26): Conv2d(2.36 M, 1.706% Params, 462.52 MMac, 2.983% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (27): ReLU(0, 0.000% Params, 100.35 KMac, 0.001% MACs, inplace=True)
    (28): Conv2d(2.36 M, 1.706% Params, 462.52 MMac, 2.983% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
    (29): ReLU(0, 0.000% Params, 100.35 KMac, 0.001% MACs, inplace=True)
    (30): MaxPool2d(0, 0.000% Params, 100.35 KMac, 0.001% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (avgpool): AdaptiveAvgPool2d(0, 0.000% Params, 25.09 KMac, 0.000% MACs, output_size=(7, 7))
  (classifier): Sequential(
    123.64 M, 89.365% Params, 123.65 MMac, 0.798% MACs,
    (0): Linear(102.76 M, 74.275% Params, 102.76 MMac, 0.663% MACs, in_features=25088, out_features=4096, bias=True)
    (1): ReLU(0, 0.000% Params, 4.1 KMac, 0.000% MACs, inplace=True)
    (2): Dropout(0, 0.000% Params, 0.0 Mac, 0.000% MACs, p=0.5, inplace=False)
    (3): Linear(16.78 M, 12.129% Params, 16.78 MMac, 0.108% MACs, in_features=4096, out_features=4096, bias=True)
    (4): ReLU(0, 0.000% Params, 4.1 KMac, 0.000% MACs, inplace=True)
    (5): Dropout(0, 0.000% Params, 0.0 Mac, 0.000% MACs, p=0.5, inplace=False)
    (6): Linear(4.1 M, 2.961% Params, 4.1 MMac, 0.026% MACs, in_features=4096, out_features=1000, bias=True)
  )
)
Computational complexity:       15.5 GMac
Number of parameters:           138.36 M
root@2c5bb9a7b477:/home/Semisup_cloudSemSeg_master#
或者使用这个代码也许:


"""
    使用 Flops 工具计算模型的计算参数量和计算量。
    author:czjing
    source: https://github.com/Lyken17/pytorch-OpCounter
    
    # pip install thop
"""

import torch

from thop import profile

from   torchvision.models import resnet50
model = resnet50()

input = torch.randn(1, 3, 224, 224)

flops,   params = profile(model, inputs=(input, ))
print('flops:', flops)
print('params:', params)

结果:

root@2c5bb9a7b477:/home/Semisup_cloudSemSeg_master# python flops.py
[INFO] Register count_convNd() for .
[INFO] Register count_bn() for .
[INFO] Register zero_ops() for .
[INFO] Register zero_ops() for .
[WARN] Cannot find rule for . Treat it as zero Macs and zero Params.
[WARN] Cannot find rule for . Treat it as zero Macs and zero Params.
[INFO] Register count_adap_avgpool() for .
[INFO] Register count_linear() for .
[WARN] Cannot find rule for . Treat it as zero Macs and zero Params.
root@2c5bb9a7b477:/home/Semisup_cloudSemSeg_master# ^C
root@2c5bb9a7b477:/home/Semisup_cloudSemSeg_master# python flops.py
[INFO] Register count_convNd() for .
[INFO] Register count_bn() for .
[INFO] Register zero_ops() for .
[INFO] Register zero_ops() for .
[WARN] Cannot find rule for . Treat it as zero Macs and zero Params.
[WARN] Cannot find rule for . Treat it as zero Macs and zero Params.
[INFO] Register count_adap_avgpool() for .
[INFO] Register count_linear() for .
[WARN] Cannot find rule for . Treat it as zero Macs and zero Params.
flops: 4111514624.0
params: 25557032.0

你可能感兴趣的:(Python,遥感图像处理,文件管理软件与技术,cnn,深度学习,pytorch)