使用pytorch计算模型的参数量和计算量

计算自带模型的例子

from torchvision.models import resnet50
import torch
import torchvision.models as models
import argparse
import vgg
import torch.nn as nn
import torch.nn.parallel
from ptflops import get_model_complexity_info

model = models.resnet50()   #调用官方的模型，

model_name = 'yolov3 cut'
flops, params = get_model_complexity_info(model, (3, 32, 32), as_strings=True, print_per_layer_stat=True)
print("%s |%s |%s" % (model_name, flops, params))

计算自己的模型：

我的模型是用GPU训练的，一个是VGG-Small的，一个是ResNet-20的。

from torchvision.models import resnet50
import torch
import torchvision.models as models
import argparse
import vgg
import torch.nn as nn
import torch.nn.parallel
from ptflops import get_model_complexity_info

# model = models.resnet50()   #调用官方的模型，
checkpoints = 'C:\\Users\\83543\\Desktop\\model_best.pth.tar'
model_names = sorted(name for name in vgg.__dict__
                     if name.islower() and not name.startswith("__")
                     and name.startswith("vgg")
                     and callable(vgg.__dict__[name]))

parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
parser.add_argument('--arch', '-a', metavar='ARCH', default='vgg_small_1w1a',
                    choices=model_names,
                    help='model architecture: ' + ' | '.join(model_names) +
                         ' (default: vgg_small)')

from ptflops import get_model_complexity_info

global args, best_prec1
args = parser.parse_args()
checkpoints = 'C:\\Users\\83543\\Desktop\\model_best.pth.tar'
pthfile = r'C:\\Users\\83543\\Desktop\\model_best.pth.tar'
model = torch.nn.DataParallel(vgg.__dict__[args.arch]())
model.cuda()
model = nn.DataParallel(model).cuda()
model.load_state_dict(torch.load("C:\\Users\\83543\\Desktop\\model_best.pth.tar")['state_dict'])

model = nn.DataParallel(model)
# model = torch.load(checkpoints)
print(model)
model1 = models.resnet50()
model_name = 'yolov3 cut'
flops, params = get_model_complexity_info(model, (3, 32, 32), as_strings=True, print_per_layer_stat=True)
print("%s |%s |%s" % (model_name, flops, params))

解释一下为什么前面有一些常量parser的代码，因为这一段是从我复现的代码里摘出来使用的，懒得去修改所以保持原样了。