pytorch如何獲得模型的計算量和參數量
方法1 自帶
pytorch自帶方法,計算模型參數總量
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))
方法2 編寫代碼
計算模型參數總量和模型計算量
def count_params(model, input_size=224): # param_sum = 0 with open('models.txt', 'w') as fm: fm.write(str(model)) # 計算模型的計算量 calc_flops(model, input_size) # 計算模型的參數總量 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)) # 計算模型的計算量 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))
方法3 thop
需要安裝thop
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")
方法4 torchstat
from torchstat import stat from torchvision.models import resnet50, resnet101, resnet152, resnext101_32x8d model = resnet50() stat(model, (3, 224, 224)) # (3,224,224)表示輸入圖片的尺寸
使用torchstat這個庫來查看網絡模型的一些信息,包括總的參數量params、MAdd、顯卡內存占用量和FLOPs等。需要安裝torchstat:
pip install torchstat
方法5 ptflops
作用:計算模型參數總量和模型計算量
安裝方法:pip install ptflops
或者
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大小,默認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() #調用官方的模型, 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))
註意,這裡輸入一定是要tuple類型,且不需要輸入batch,直接輸入輸入通道數量與尺寸,如(3,320,320) 320為網絡輸入尺寸。
輸出為網絡模型的總參數量(單位M,即百萬)與計算量(單位G,即十億)
方法6 torchsummary
安裝:pip install torchsummary
使用方法:
from torchsummary import summary ... summary(your_model, input_size=(channels, H, W))
作用:
1、每一層的類型、shape 和 參數量
2、模型整體的參數量
3、模型大小,和 fp/bp 一次需要的內存大小,可以用來估計最佳 batch_size
補充:pytorch計算模型算力與參數大小
ptflops介紹
官方鏈接
這個腳本設計用於計算卷積神經網絡中乘法-加法操作的理論數量。它還可以計算參數的數量和打印給定網絡的每層計算成本。
支持layer:Conv1d/2d/3d,ConvTranspose2d,BatchNorm1d/2d/3d,激活(ReLU, PReLU, ELU, ReLU6, LeakyReLU),Linear,Upsample,Poolings (AvgPool1d/2d/3d、MaxPool1d/2d/3d、adaptive ones)
安裝要求:Pytorch >= 0.4.1, torchvision >= 0.2.1
get_model_complexity_info()
get_model_complexity_info是ptflops下的一個方法,可以計算出網絡的算力與模型參數大小,並且可以輸出每層的算力消耗。
栗子
以輸出Mobilenet_v2算力信息為例:
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)
從圖中可以看到,MobileNetV2在輸入圖像尺寸為(3, 224, 224)的情況下將會產生3.505MB的參數,算力消耗為0.32G,同時還打印出瞭每個層所占用的算力,權重參數數量。當然,整個模型的算力大小與模型大小也被存到瞭變量ops與params中。
以上為個人經驗,希望能給大傢一個參考,也希望大傢多多支持WalkonNet。
推薦閱讀:
- pytorch查看網絡參數顯存占用量等操作
- pytorch 計算Parameter和FLOP的操作
- Python LeNet網絡詳解及pytorch實現
- Pytorch實現ResNet網絡之Residual Block殘差塊
- pytorch lstm gru rnn 得到每個state輸出的操作