pytorch如何獲得模型的計(jì)算量和參數(shù)量

total = sum([param.nelement() for param in model.parameters()])
print("Number of parameter: %.2fM" % (total/1e6))

total = sum(p.numel() for p in model.parameters())
print("Total params: %.2fM" % (total/1e6))

def count_params(model, input_size=224):
    # param_sum = 0
    with open('models.txt', 'w') as fm:
        fm.write(str(model))
 
    # 計(jì)算模型的計(jì)算量
    calc_flops(model, input_size)
 
    # 計(jì)算模型的參數(shù)總量
    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
    params = sum([np.prod(p.size()) for p in model_parameters])
 
    print('The network has {} params.'.format(params)) 
 
# 計(jì)算模型的計(jì)算量
def calc_flops(model, input_size):
 
    def conv_hook(self, input, output):
        batch_size, input_channels, input_height, input_width = input[0].size()
        output_channels, output_height, output_width = output[0].size()
 
        kernel_ops = self.kernel_size[0] * self.kernel_size[1] * (self.in_channels / self.groups) * (
            2 if multiply_adds else 1)
        bias_ops = 1 if self.bias is not None else 0
 
        params = output_channels * (kernel_ops + bias_ops)
        flops = batch_size * params * output_height * output_width
 
        list_conv.append(flops)
 
    def linear_hook(self, input, output):
        batch_size = input[0].size(0) if input[0].dim() == 2 else 1
 
        weight_ops = self.weight.nelement() * (2 if multiply_adds else 1)
        bias_ops = self.bias.nelement()
 
        flops = batch_size * (weight_ops + bias_ops)
        list_linear.append(flops)
 
    def bn_hook(self, input, output):
        list_bn.append(input[0].nelement())
 
    def relu_hook(self, input, output):
        list_relu.append(input[0].nelement())
 
    def pooling_hook(self, input, output):
        batch_size, input_channels, input_height, input_width = input[0].size()
        output_channels, output_height, output_width = output[0].size()
 
        kernel_ops = self.kernel_size * self.kernel_size
        bias_ops = 0
        params = output_channels * (kernel_ops + bias_ops)
        flops = batch_size * params * output_height * output_width
 
        list_pooling.append(flops)
 
    def foo(net):
        childrens = list(net.children())
        if not childrens:
            if isinstance(net, torch.nn.Conv2d):
                net.register_forward_hook(conv_hook)
            if isinstance(net, torch.nn.Linear):
                net.register_forward_hook(linear_hook)
            if isinstance(net, torch.nn.BatchNorm2d):
                net.register_forward_hook(bn_hook)
            if isinstance(net, torch.nn.ReLU):
                net.register_forward_hook(relu_hook)
            if isinstance(net, torch.nn.MaxPool2d) or isinstance(net, torch.nn.AvgPool2d):
                net.register_forward_hook(pooling_hook)
            return
        for c in childrens:
            foo(c)
 
    multiply_adds = False
    list_conv, list_bn, list_relu, list_linear, list_pooling = [], [], [], [], []
    foo(model)
    if '0.4.' in torch.__version__:
        if assets.USE_GPU:
            input = torch.cuda.FloatTensor(torch.rand(2, 3, input_size, input_size).cuda())
        else:
            input = torch.FloatTensor(torch.rand(2, 3, input_size, input_size))
    else:
        input = Variable(torch.rand(2, 3, input_size, input_size), requires_grad=True)
    _ = model(input)
 
    total_flops = (sum(list_conv) + sum(list_linear) + sum(list_bn) + sum(list_relu) + sum(list_pooling))
 
    print('  + Number of FLOPs: %.2fM' % (total_flops / 1e6 / 2))

pip install thop

from thop import profile 
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input,))
print(flops)
print(params)

from torchvision.models import resnet50
from thop import profile
 
# model = resnet50()
checkpoints = '模型path'
model = torch.load(checkpoints)
model_name = 'yolov3 cut asff'
input = torch.randn(1, 3, 224, 224)
flops, params = profile(model, inputs=(input, ),verbose=True)
print("%s | %.2f | %.2f" % (model_name, params / (1000 ** 2), flops / (1000 ** 3)))#這里除以1000的平方，是為了化成M的單位，

from thop import clever_format
macs, params = clever_format([flops, params], "%.3f")

from torchstat import stat
from torchvision.models import resnet50, resnet101, resnet152, resnext101_32x8d
 
model = resnet50()
stat(model, (3, 224, 224))  #  (3,224,224)表示輸入圖片的尺寸

pip install torchstat

pip install --upgrade git+https://github.com/sovrasov/flops-counter.pytorch.git

import torchvision.models as models
import torch
from ptflops import get_model_complexity_info
with torch.cuda.device(0):
  net = models.resnet18()
  flops, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, print_per_layer_stat=True) #不用寫batch_size大小，默認(rèn)batch_size=1
  print('Flops:  ' + flops)
  print('Params: ' + params)

from torchvision.models import resnet50
import torch
import torchvision.models as models
# import torch
from ptflops import get_model_complexity_info
 
# model = models.resnet50()   #調(diào)用官方的模型，
checkpoints = '自己模型的path'
model = torch.load(checkpoints)
model_name = 'yolov3 cut'
flops, params = get_model_complexity_info(model, (3,320,320),as_strings=True,print_per_layer_stat=True)
print("%s |%s |%s" % (model_name,flops,params))

from torchsummary import summary
...
summary(your_model, input_size=(channels, H, W))

from ptflops import get_model_complexity_info
from torchvision import models
net = models.mobilenet_v2()
ops, params = get_model_complexity_info(net, (3, 224, 224), as_strings=True, 										print_per_layer_stat=True, verbose=True)