计算自带模型的例子
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的代码,因为这一段是从我复现的代码里摘出来使用的,懒得去修改所以保持原样了。