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中。