5种方法获取Torch网络模型参数量计算量等信息
目录
-
- 模型参数量和计算量是什么
- 为什么要统计模型参数量和计算量
- 常见的模型参数量和计算量的计算方法
- 代码
-
-
- 方法1 利用torch自带方法，自写函数
- 方法2 torchsummary.summary
- 方法3 torchstat.stat
- 方法4 thop.profile
- 方法5 ptflops.get_model_complexity_info
-
模型参数量和计算量是什么
- 计算量是指网络模型需要计算的运算次数，参数量是指网络模型自带的参数数量多少
- 计算量对应时间复杂度，参数量对应于空间复杂度
- 计算量决定了网络执行时间的长短，参数量决定了占用显存的量
为什么要统计模型参数量和计算量
- 好的网络模型不仅要求精度准，还要要求模型的参数量和计算量不大，才能有利于部署
- 统计模型的参数量和计算量可以用于不同网络模型之间的对比分析
- 有的模型虽然参数量相同，但是可能因为连接方式和结构等不同而导致计算量不同
常见的模型参数量和计算量的计算方法
- 卷积层
Time ∼ O ( ∑ l = 1 D M l 2 ⋅ K l 2 ⋅ C l − 1 ⋅ C l ) Space ∼ O ( ∑ l = 1 D K l 2 ⋅ C l − 1 ⋅ C l + ∑ l = 1 D M 2 ⋅ C l ) \begin{array}{r} \text { Time } \sim O\left(\sum_{l=1}^{D} M_{l}^{2} \cdot K_{l}^{2} \cdot C_{l-1} \cdot C_{l}\right) \\ \text { Space } \sim O\left(\sum_{l=1}^{D} K_{l}^{2} \cdot C_{l-1} \cdot C_{l}+\sum_{l=1}^{D} M^{2} \cdot C_{l}\right) \end{array} Time ∼O(∑l=1DMl2⋅Kl2⋅Cl−1⋅Cl) Space ∼O(∑l=1DKl2⋅Cl−1⋅Cl+∑l=1DM2⋅Cl)
其中， K K K表示核尺寸， C l C_l Cl代表数量第 l l l层通道数
- 全连接层
计 算 量 = 参 数 量 = w e i g h t i n × w e i g h t o u t 计算量=参数量=weight_{in}×weight_{out} 计算量=参数量=weightin×weightout
代码
定义网络与输入，需安装torch和torchvision包
import torch.nn as nn
import torch
from torchvision.models import vgg16
net = vgg16()
x = torch.rand(1,3,224,224)
方法1 利用torch自带方法，自写函数
自定义函数，参考自 https://blog.csdn.net/qq_33757398/article/details/109210240
'''方法1，自定义函数 参考自 https://blog.csdn.net/qq_33757398/article/details/109210240'''
def model_structure(model):
blank = ' '
print('-' * 90)
print('|' + ' ' * 11 + 'weight name' + ' ' * 10 + '|' \
+ ' ' * 15 + 'weight shape' + ' ' * 15 + '|' \
+ ' ' * 3 + 'number' + ' ' * 3 + '|')
print('-' * 90)
num_para = 0
type_size = 1 # 如果是浮点数就是4
for index, (key, w_variable) in enumerate(model.named_parameters()):
if len(key) <= 30:
key = key + (30 - len(key)) * blank
shape = str(w_variable.shape)
if len(shape) <= 40:
shape = shape + (40 - len(shape)) * blank
each_para = 1
for k in w_variable.shape:
each_para *= k
num_para += each_para
str_num = str(each_para)
if len(str_num) <= 10:
str_num = str_num + (10 - len(str_num)) * blank
print('| {} | {} | {} |'.format(key, shape, str_num))
print('-' * 90)
print('The total number of parameters: ' + str(num_para))
print('The parameters of Model {}: {:4f}M'.format(model._get_name(), num_para * type_size / 1000 / 1000))
print('-' * 90)
model_structure(net)
输出：
------------------------------------------------------------------------------------------
| weight name | weight shape | number |
------------------------------------------------------------------------------------------
| features.0.weight | torch.Size([64, 3, 3, 3]) | 1728 |
| features.0.bias | torch.Size([64]) | 64 |
| features.2.weight | torch.Size([64, 64, 3, 3]) | 36864 |
| features.2.bias | torch.Size([64]) | 64 |
| features.5.weight | torch.Size([128, 64, 3, 3]) | 73728 |
| features.5.bias | torch.Size([128]) | 128 |
| features.7.weight | torch.Size([128, 128, 3, 3]) | 147456 |
| features.7.bias | torch.Size([128]) | 128 |
| features.10.weight | torch.Size([256, 128, 3, 3]) | 294912 |
| features.10.bias | torch.Size([256]) | 256 |
| features.12.weight | torch.Size([256, 256, 3, 3]) | 589824 |
| features.12.bias | torch.Size([256]) | 256 |
| features.14.weight | torch.Size([256, 256, 3, 3]) | 589824 |
| features.14.bias | torch.Size([256]) | 256 |
| features.17.weight | torch.Size([512, 256, 3, 3]) | 1179648 |
| features.17.bias | torch.Size([512]) | 512 |
| features.19.weight | torch.Size([512, 512, 3, 3]) | 2359296 |
| features.19.bias | torch.Size([512]) | 512 |
| features.21.weight | torch.Size([512, 512, 3, 3]) | 2359296 |
| features.21.bias | torch.Size([512]) | 512 |
| features.24.weight | torch.Size([512, 512, 3, 3]) | 2359296 |
| features.24.bias | torch.Size([512]) | 512 |
| features.26.weight | torch.Size([512, 512, 3, 3]) | 2359296 |
| features.26.bias | torch.Size([512]) | 512 |
| features.28.weight | torch.Size([512, 512, 3, 3]) | 2359296 |
| features.28.bias | torch.Size([512]) | 512 |
| classifier.0.weight | torch.Size([4096, 25088]) | 102760448 |
| classifier.0.bias | torch.Size([4096]) | 4096 |
| classifier.3.weight | torch.Size([4096, 4096]) | 16777216 |
| classifier.3.bias | torch.Size([4096]) | 4096 |
| classifier.6.weight | torch.Size([1000, 4096]) | 4096000 |
| classifier.6.bias | torch.Size([1000]) | 1000 |
------------------------------------------------------------------------------------------
The total number of parameters: 138357544
The parameters of Model VGG: 138.357544M
------------------------------------------------------------------------------------------
方法2 torchsummary.summary
from torchsummary import summary
summary(net,input_size=(3,224,224))
输出：
----------------------------------------------------------------
Layer (type) Output Shape Param #
================================================================
Conv2d-1 [-1, 64, 224, 224] 1,792
ReLU-2 [-1, 64, 224, 224] 0
Conv2d-3 [-1, 64, 224, 224] 36,928
ReLU-4 [-1, 64, 224, 224] 0
MaxPool2d-5 [-1, 64, 112, 112] 0
Conv2d-6 [-1, 128, 112, 112] 73,856
ReLU-7 [-1, 128, 112, 112] 0
Conv2d-8 [-1, 128, 112, 112] 147,584
ReLU-9 [-1, 128, 112, 112] 0
MaxPool2d-10 [-1, 128, 56, 56] 0
Conv2d-11 [-1, 256, 56, 56] 295,168
ReLU-12 [-1, 256, 56, 56] 0
Conv2d-13 [-1, 256, 56, 56] 590,080
ReLU-14 [-1, 256, 56, 56] 0
Conv2d-15 [-1, 256, 56, 56] 590,080
ReLU-16 [-1, 256, 56, 56] 0
MaxPool2d-17 [-1, 256, 28, 28] 0
Conv2d-18 [-1, 512, 28, 28] 1,180,160
ReLU-19 [-1, 512, 28, 28] 0
Conv2d-20 [-1, 512, 28, 28] 2,359,808
ReLU-21 [-1, 512, 28, 28] 0
Conv2d-22 [-1, 512, 28, 28] 2,359,808
ReLU-23 [-1, 512, 28, 28] 0
MaxPool2d-24 [-1, 512, 14, 14] 0
Conv2d-25 [-1, 512, 14, 14] 2,359,808
ReLU-26 [-1, 512, 14, 14] 0
Conv2d-27 [-1, 512, 14, 14] 2,359,808
ReLU-28 [-1, 512, 14, 14] 0
Conv2d-29 [-1, 512, 14, 14] 2,359,808
ReLU-30 [-1, 512, 14, 14] 0
MaxPool2d-31 [-1, 512, 7, 7] 0
AdaptiveAvgPool2d-32 [-1, 512, 7, 7] 0
Linear-33 [-1, 4096] 102,764,544
ReLU-34 [-1, 4096] 0
Dropout-35 [-1, 4096] 0
Linear-36 [-1, 4096] 16,781,312
ReLU-37 [-1, 4096] 0
Dropout-38 [-1, 4096] 0
Linear-39 [-1, 1000] 4,097,000
================================================================
Total params: 138,357,544
Trainable params: 138,357,544
Non-trainable params: 0
----------------------------------------------------------------
Input size (MB): 0.57
Forward/backward pass size (MB): 218.78
Params size (MB): 527.79
Estimated Total Size (MB): 747.15
----------------------------------------------------------------
方法3 torchstat.stat
from torchstat import stat
stat(net,(3,224,224))
输出：
[MAdd]: AdaptiveAvgPool2d is not supported!
[Flops]: AdaptiveAvgPool2d is not supported!
[Memory]: AdaptiveAvgPool2d is not supported!
[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!
[MAdd]: Dropout is not supported!
[Flops]: Dropout is not supported!
[Memory]: Dropout is not supported!
module name input shape output shape params memory(MB) MAdd Flops MemRead(B) MemWrite(B) duration[%] MemR+W(B)
0 features.0 3 224 224 64 224 224 1792.0 12.25 173,408,256.0 89,915,392.0 609280.0 12845056.0 4.91% 13454336.0
1 features.1 64 224 224 64 224 224 0.0 12.25 3,211,264.0 3,211,264.0 12845056.0 12845056.0 1.23% 25690112.0
2 features.2 64 224 224 64 224 224 36928.0 12.25 3,699,376,128.0 1,852,899,328.0 12992768.0 12845056.0 10.66% 25837824.0
3 features.3 64 224 224 64 224 224 0.0 12.25 3,211,264.0 3,211,264.0 12845056.0 12845056.0 1.23% 25690112.0
4 features.4 64 224 224 64 112 112 0.0 3.06 2,408,448.0 3,211,264.0 12845056.0 3211264.0 6.97% 16056320.0
5 features.5 64 112 112 128 112 112 73856.0 6.12 1,849,688,064.0 926,449,664.0 3506688.0 6422528.0 4.92% 9929216.0
6 features.6 128 112 112 128 112 112 0.0 6.12 1,605,632.0 1,605,632.0 6422528.0 6422528.0 0.41% 12845056.0
7 features.7 128 112 112 128 112 112 147584.0 6.12 3,699,376,128.0 1,851,293,696.0 7012864.0 6422528.0 7.79% 13435392.0
8 features.8 128 112 112 128 112 112 0.0 6.12 1,605,632.0 1,605,632.0 6422528.0 6422528.0 0.41% 12845056.0
9 features.9 128 112 112 128 56 56 0.0 1.53 1,204,224.0 1,605,632.0 6422528.0 1605632.0 3.28% 8028160.0
10 features.10 128 56 56 256 56 56 295168.0 3.06 1,849,688,064.0 925,646,848.0 2786304.0 3211264.0 3.69% 5997568.0
11 features.11 256 56 56 256 56 56 0.0 3.06 802,816.0 802,816.0 3211264.0 3211264.0 0.00% 6422528.0
12 features.12 256 56 56 256 56 56 590080.0 3.06 3,699,376,128.0 1,850,490,880.0 5571584.0 3211264.0 6.56% 8782848.0
13 features.13 256 56 56 256 56 56 0.0 3.06 802,816.0 802,816.0 3211264.0 3211264.0 0.41% 6422528.0
14 features.14 256 56 56 256 56 56 590080.0 3.06 3,699,376,128.0 1,850,490,880.0 5571584.0 3211264.0 6.56% 8782848.0
15 features.15 256 56 56 256 56 56 0.0 3.06 802,816.0 802,816.0 3211264.0 3211264.0 0.00% 6422528.0
16 features.16 256 56 56 256 28 28 0.0 0.77 602,112.0 802,816.0 3211264.0 802816.0 2.46% 4014080.0
17 features.17 256 28 28 512 28 28 1180160.0 1.53 1,849,688,064.0 925,245,440.0 5523456.0 1605632.0 4.51% 7129088.0
18 features.18 512 28 28 512 28 28 0.0 1.53 401,408.0 401,408.0 1605632.0 1605632.0 0.00% 3211264.0
19 features.19 512 28 28 512 28 28 2359808.0 1.53 3,699,376,128.0 1,850,089,472.0 11044864.0 1605632.0 7.38% 12650496.0
20 features.20 512 28 28 512 28 28 0.0 1.53 401,408.0 401,408.0 1605632.0 1605632.0 0.00% 3211264.0
21 features.21 512 28 28 512 28 28 2359808.0 1.53 3,699,376,128.0 1,850,089,472.0 11044864.0 1605632.0 7.38% 12650496.0
22 features.22 512 28 28 512 28 28 0.0 1.53 401,408.0 401,408.0 1605632.0 1605632.0 0.00% 3211264.0
23 features.23 512 28 28 512 14 14 0.0 0.38 301,056.0 401,408.0 1605632.0 401408.0 0.82% 2007040.0
24 features.24 512 14 14 512 14 14 2359808.0 0.38 924,844,032.0 462,522,368.0 9840640.0 401408.0 3.28% 10242048.0
25 features.25 512 14 14 512 14 14 0.0 0.38 100,352.0 100,352.0 401408.0 401408.0 0.00% 802816.0
26 features.26 512 14 14 512 14 14 2359808.0 0.38 924,844,032.0 462,522,368.0 9840640.0 401408.0 2.87% 10242048.0
27 features.27 512 14 14 512 14 14 0.0 0.38 100,352.0 100,352.0 401408.0 401408.0 0.00% 802816.0
28 features.28 512 14 14 512 14 14 2359808.0 0.38 924,844,032.0 462,522,368.0 9840640.0 401408.0 2.46% 10242048.0
29 features.29 512 14 14 512 14 14 0.0 0.38 100,352.0 100,352.0 401408.0 401408.0 0.00% 802816.0
30 features.30 512 14 14 512 7 7 0.0 0.10 75,264.0 100,352.0 401408.0 100352.0 0.00% 501760.0
31 avgpool 512 7 7 512 7 7 0.0 0.10 0.0 0.0 0.0 0.0 0.41% 0.0
32 classifier.0 25088 4096 102764544.0 0.02 205,516,800.0 102,760,448.0 411158528.0 16384.0 7.79% 411174912.0
33 classifier.1 4096 4096 0.0 0.02 4,096.0 4,096.0 16384.0 16384.0 0.00% 32768.0
34 classifier.2 4096 4096 0.0 0.02 0.0 0.0 0.0 0.0 0.00% 0.0
35 classifier.3 4096 4096 16781312.0 0.02 33,550,336.0 16,777,216.0 67141632.0 16384.0 1.23% 67158016.0
36 classifier.4 4096 4096 0.0 0.02 4,096.0 4,096.0 16384.0 16384.0 0.00% 32768.0
37 classifier.5 4096 4096 0.0 0.02 0.0 0.0 0.0 0.0 0.00% 0.0
38 classifier.6 4096 1000 4097000.0 0.00 8,191,000.0 4,096,000.0 16404384.0 4000.0 0.41% 16408384.0
total 138357544.0 109.39 30,958,666,264.0 15,503,489,024.0 16404384.0 4000.0 100.00% 783170624.0
============================================================================================================================================================
Total params: 138,357,544
------------------------------------------------------------------------------------------------------------------------------------------------------------
Total memory: 109.39MB
Total MAdd: 30.96GMAdd
Total Flops: 15.5GFlops
Total MemR+W: 746.89MB
方法4 thop.profile
from thop import profile
from thop import clever_format
flops, params = profile(net, inputs=(x, ))
# print(flops, params)
macs, params = clever_format([flops, params], "%.3f")
print(macs,params)
输出：
15.484G 138.358M
方法5 ptflops.get_model_complexity_info
from ptflops import get_model_complexity_info
flops, params = get_model_complexity_info(net, (3,224,224),as_strings=True,print_per_layer_stat=True)
print("%s %s" % (flops,params))
输出
VGG(
138.358 M, 100.000% Params, 15.504 GMac, 100.000% MACs,
(features): Sequential(
14.715 M, 10.635% Params, 15.38 GMac, 99.202% MACs,
(0): Conv2d(0.002 M, 0.001% Params, 0.09 GMac, 0.580% MACs, 3, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU(0.0 M, 0.000% Params, 0.003 GMac, 0.021% MACs, inplace=True)
(2): Conv2d(0.037 M, 0.027% Params, 1.853 GMac, 11.951% MACs, 64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(3): ReLU(0.0 M, 0.000% Params, 0.003 GMac, 0.021% MACs, inplace=True)
(4): MaxPool2d(0.0 M, 0.000% Params, 0.003 GMac, 0.021% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(5): Conv2d(0.074 M, 0.053% Params, 0.926 GMac, 5.976% MACs, 64, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(6): ReLU(0.0 M, 0.000% Params, 0.002 GMac, 0.010% MACs, inplace=True)
(7): Conv2d(0.148 M, 0.107% Params, 1.851 GMac, 11.941% MACs, 128, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(8): ReLU(0.0 M, 0.000% Params, 0.002 GMac, 0.010% MACs, inplace=True)
(9): MaxPool2d(0.0 M, 0.000% Params, 0.002 GMac, 0.010% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(10): Conv2d(0.295 M, 0.213% Params, 0.926 GMac, 5.971% MACs, 128, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(11): ReLU(0.0 M, 0.000% Params, 0.001 GMac, 0.005% MACs, inplace=True)
(12): Conv2d(0.59 M, 0.426% Params, 1.85 GMac, 11.936% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(13): ReLU(0.0 M, 0.000% Params, 0.001 GMac, 0.005% MACs, inplace=True)
(14): Conv2d(0.59 M, 0.426% Params, 1.85 GMac, 11.936% MACs, 256, 256, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(15): ReLU(0.0 M, 0.000% Params, 0.001 GMac, 0.005% MACs, inplace=True)
(16): MaxPool2d(0.0 M, 0.000% Params, 0.001 GMac, 0.005% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(17): Conv2d(1.18 M, 0.853% Params, 0.925 GMac, 5.968% MACs, 256, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(18): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.003% MACs, inplace=True)
(19): Conv2d(2.36 M, 1.706% Params, 1.85 GMac, 11.933% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(20): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.003% MACs, inplace=True)
(21): Conv2d(2.36 M, 1.706% Params, 1.85 GMac, 11.933% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(22): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.003% MACs, inplace=True)
(23): MaxPool2d(0.0 M, 0.000% Params, 0.0 GMac, 0.003% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(24): Conv2d(2.36 M, 1.706% Params, 0.463 GMac, 2.983% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(25): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.001% MACs, inplace=True)
(26): Conv2d(2.36 M, 1.706% Params, 0.463 GMac, 2.983% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(27): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.001% MACs, inplace=True)
(28): Conv2d(2.36 M, 1.706% Params, 0.463 GMac, 2.983% MACs, 512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(29): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.001% MACs, inplace=True)
(30): MaxPool2d(0.0 M, 0.000% Params, 0.0 GMac, 0.001% MACs, kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
)
(avgpool): AdaptiveAvgPool2d(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, output_size=(7, 7))
(classifier): Sequential(
123.643 M, 89.365% Params, 0.124 GMac, 0.798% MACs,
(0): Linear(102.765 M, 74.275% Params, 0.103 GMac, 0.663% MACs, in_features=25088, out_features=4096, bias=True)
(1): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, inplace=True)
(2): Dropout(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, p=0.5, inplace=False)
(3): Linear(16.781 M, 12.129% Params, 0.017 GMac, 0.108% MACs, in_features=4096, out_features=4096, bias=True)
(4): ReLU(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, inplace=True)
(5): Dropout(0.0 M, 0.000% Params, 0.0 GMac, 0.000% MACs, p=0.5, inplace=False)
(6): Linear(4.097 M, 2.961% Params, 0.004 GMac, 0.026% MACs, in_features=4096, out_features=1000, bias=True)
)
)
15.5 GMac 138.36 M
个人更喜欢方法3