之前寫過一篇手動計算卷積神經(jīng)網(wǎng)絡(luò)參數(shù)量和FLOPs的計算公式的博客：卷積神經(jīng)網(wǎng)絡(luò)模型參數(shù)量和運算量計算方法俄认。
現(xiàn)介紹一款自動計算模型參數(shù)量、FLOPs洪乍、乘加數(shù)以及所需內(nèi)存等數(shù)據(jù)的輪子：torchstat眯杏。
食用方法如下：
1.安裝：
pip install torchstat
2.使用：
In [3]: import torch
In [4]: from torchstat import stat
In [5]: import torchvision.models as models
In [6]: net = models.vgg11()
In [7]: stat(net,(3,224,224)) # (3,224,224)表示輸入圖片的尺寸
運行結(jié)果：
[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 14.24% 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.40% 25690112.0
2 features.2 64 224 224 64 112 112 0.0 3.06 2,408,448.0 3,211,264.0 12845056.0 3211264.0 11.61% 16056320.0
3 features.3 64 112 112 128 112 112 73856.0 6.12 1,849,688,064.0 926,449,664.0 3506688.0 6422528.0 5.33% 9929216.0
4 features.4 128 112 112 128 112 112 0.0 6.12 1,605,632.0 1,605,632.0 6422528.0 6422528.0 0.65% 12845056.0
5 features.5 128 112 112 128 56 56 0.0 1.53 1,204,224.0 1,605,632.0 6422528.0 1605632.0 4.59% 8028160.0
6 features.6 128 56 56 256 56 56 295168.0 3.06 1,849,688,064.0 925,646,848.0 2786304.0 3211264.0 4.42% 5997568.0
7 features.7 256 56 56 256 56 56 0.0 3.06 802,816.0 802,816.0 3211264.0 3211264.0 0.29% 6422528.0
8 features.8 256 56 56 256 56 56 590080.0 3.06 3,699,376,128.0 1,850,490,880.0 5571584.0 3211264.0 8.13% 8782848.0
9 features.9 256 56 56 256 56 56 0.0 3.06 802,816.0 802,816.0 3211264.0 3211264.0 0.25% 6422528.0
10 features.10 256 56 56 256 28 28 0.0 0.77 602,112.0 802,816.0 3211264.0 802816.0 2.19% 4014080.0
11 features.11 256 28 28 512 28 28 1180160.0 1.53 1,849,688,064.0 925,245,440.0 5523456.0 1605632.0 4.72% 7129088.0
12 features.12 512 28 28 512 28 28 0.0 1.53 401,408.0 401,408.0 1605632.0 1605632.0 0.10% 3211264.0
13 features.13 512 28 28 512 28 28 2359808.0 1.53 3,699,376,128.0 1,850,089,472.0 11044864.0 1605632.0 8.77% 12650496.0
14 features.14 512 28 28 512 28 28 0.0 1.53 401,408.0 401,408.0 1605632.0 1605632.0 0.09% 3211264.0
15 features.15 512 28 28 512 14 14 0.0 0.38 301,056.0 401,408.0 1605632.0 401408.0 1.14% 2007040.0
16 features.16 512 14 14 512 14 14 2359808.0 0.38 924,844,032.0 462,522,368.0 9840640.0 401408.0 3.53% 10242048.0
17 features.17 512 14 14 512 14 14 0.0 0.38 100,352.0 100,352.0 401408.0 401408.0 0.03% 802816.0
18 features.18 512 14 14 512 14 14 2359808.0 0.38 924,844,032.0 462,522,368.0 9840640.0 401408.0 3.64% 10242048.0
19 features.19 512 14 14 512 14 14 0.0 0.38 100,352.0 100,352.0 401408.0 401408.0 0.03% 802816.0
20 features.20 512 14 14 512 7 7 0.0 0.10 75,264.0 100,352.0 401408.0 100352.0 0.32% 501760.0
21 avgpool 512 7 7 512 7 7 0.0 0.10 0.0 0.0 0.0 0.0 0.25% 0.0
22 classifier.0 25088 4096 102764544.0 0.02 205,516,800.0 102,760,448.0 411158528.0 16384.0 20.16% 411174912.0
23 classifier.1 4096 4096 0.0 0.02 4,096.0 4,096.0 16384.0 16384.0 0.02% 32768.0
24 classifier.2 4096 4096 0.0 0.02 0.0 0.0 0.0 0.0 0.03% 0.0
25 classifier.3 4096 4096 16781312.0 0.02 33,550,336.0 16,777,216.0 67141632.0 16384.0 3.26% 67158016.0
26 classifier.4 4096 4096 0.0 0.02 4,096.0 4,096.0 16384.0 16384.0 0.01% 32768.0
27 classifier.5 4096 4096 0.0 0.02 0.0 0.0 0.0 0.0 0.01% 0.0
28 classifier.6 4096 1000 4097000.0 0.00 8,191,000.0 4,096,000.0 16404384.0 4000.0 0.79% 16408384.0
total 132863336.0 62.69 15,230,196,248.0 7,630,071,808.0 16404384.0 4000.0 100.00% 663250240.0
===========================================================================================================================================================
Total params: 132,863,336
-----------------------------------------------------------------------------------------------------------------------------------------------------------
Total memory: 62.69MB
Total MAdd: 15.23GMAdd
Total Flops: 7.63GFlops
Total MemR+W: 632.52MB
支持環(huán)境：
- Python 3.6+
- Pytorch 0.4.0+
- Pandas 0.23.4+
- NumPy 1.14.3+
