模型的显存和参数量计算

写在前面：以此记录关于模型显存和参数量的一些理解和计算。

首先是“运算量”和“参数量”两个概念：
参数量：这个比较好理解，例如卷积层中的卷积核c_i*k*k*n_o，其参数量就是相乘的结果。而且，无论输入图像的尺寸怎么变（YOLO实现中的multi scale训练策略），只要模型结构确定，参数量就固定。还需要注意，参数都是FP32（4字节）存放，所以模型大小是参数量*4。
运算量：使用FLOPs衡量，代表浮点运算次数，这个可以衡量算法/模型的复杂度。

还要区分两个容易混淆的量：FLOPS和FLOPs：
FLOPS：Floating point Opreations Per Second，每秒浮点运算次数，理解为计算速度，是一个衡量硬件的标准。GPU算力描述的就是这个，这些数值的单位为MM=10^12次。
FLOPs：Floating point Opreations，s是复数，为浮点数操作数。理解为算法的计算量。paper中通常使用的是GFLOPs，即10亿次浮点运算。

一、显存

参数的显存占用：只有有参数的层，才会有显存占用。这部分的显存占用和输入无关，模型加载完成之后就会占用。

有参数的层主要包括：卷积、全连接、BatchNorm、Embedding等等

无参数的层主要包括：多数的激活层(Sigmoid/ReLU)、池化层、Dropout等等

具体来说，模型参数数目(不考虑偏置b)：

• Linear(M->N)：M×N
• Conv2d(Cin，Cout，K)：Cin×Cout×k×k
• BatchNorm(N)：2N
• Embedding(N，W)：N×W

参数占用显存=参数数目×n

n=4：float32

n=2：float16

n=8：double64

梯度与动量的显存占用

1. 优化器如果是SGD：

   

   可以看出，除了保存W之外还要保存对应的梯度

   

   因此显存占用等于参数从占用的显存×2

2. 优化器如果是带Momentum-SGD

这时候还需要保存动量，因此显存×3

如果是Adam优化器，动量占用的显存更多，显存×4

故模型中与输入输出无关的显存占用包括：参数W、梯度dW(一般与参数一样)、优化器的动量

输入输出显存占用

主要看输出feature map的形状

输入可以看成上一层的输出

因此模型输出显存占用为：

1. 每一层feature map的形状(多维数组的形状)
2. 保存输出对应的梯度泳衣反向传播(链式法则)
3. 显存占用与batch size成正比
4. 模型输出不需要存储相应的动量信息

因此，总的显存公式为：

显存占用=模型显存(参数)+batch_size×每个样本显存(输出和梯度动量)

节省显存的方法：

1. 降低batch_size
2. 下采样(NCHW->(1/4)*NCHW)
3. 减少全连接层(一般只留最后一层分类用的全连接层)

二、计算量

卷积神经网络的参数位于卷积层和全连接层。

三、举例

以VGG16为例

各层的具体情况如下：

网络层（操作）	输入	filter	stride	padding	输出	参数量计算公式	参数量	显存占用
Input	224x224x3				224x224x3		0	224x224x3
Conv3-64	224x224x3	3x3x64	1	1	224x224x64	3x3x3x64 (+ 64)(括号里面是bias项）	1728(1792)	224x224x64
Conv3-64	224x224x64	3x3x64	1	1	224x224x64	3x3x64x64 (+ 64)	36864(36928)	224x224x64
MaxPool2	224x224x64	2x2	2	0	112x112x64		0	112x112x64
Conv3-128	112x112x64	3x3x128	1	1	112x112x128	3x3x64x128 (+ 128)	73728(73856)	112x112x128
Conv3-128	112x112x128	3x3x128	1	1	112x112x128	3x3x128x128 (+ 128)	147456(147584)	112x112x128
MaxPool2	112x112x128	2x2	2	0	56x56x128		0	56x56x128
Conv3-256	56x56x128	3x3x256	1	1	56x56x256	3x3x128x256 (+ 256)	294912(295168)	56x56x256
Conv3-256	56x56x256	3x3x256	1	1	56x56x256	3x3x256x256 (+ 256)	589824(590080)	56x56x256
Conv3-256	56x56x256	3x3x256	1	1	56x56x256	3x3x256x256 (+ 256)	589824(590080)	56x56x256
MaxPool2	56x56x256	2x2	2	0	28x28x256		0	28x28x256
Conv3-512	28x28x256	3x3x512	1	1	28x28x512	3x3x256x512 (+ 512)	1179648(1180160)	28x28x512
Conv3-512	28x28x512	3x3x512	1	1	28x28x512	3x3x512x512 (+ 512)	2359296(2359808)	28x28x512
Conv3-512	28x28x512	3x3x512	1	1	28x28x512	3x3x512x512 (+ 512)	2359296(2359808)	28x28x512
MaxPool2	28x28x512	2x2	2	0	14x14x512		0	14x14x512
Conv3-512	14x14x512	3x3x512	1	1	14x14x512	3x3x512x512 (+ 512)	2359296(2359808)	14x14x512
Conv3-512	14x14x512	3x3x512	1	1	14x14x512	3x3x512x512 (+ 512)	2359296(2359808)	14x14x512
Conv3-512	14x14x512	3x3x512	1	1	14x14x512	3x3x512x512 (+ 512)	2359296(2359808)	14x14x512
MaxPool2	14x14x512	2x2	2	0	7x7x512		0	7x7x512
FC1	7x7x512				4096	7x7x512x4096 (+ 4096)	102760448(102764544)	4096
FC2	4096				4096	4096*4096 (+ 4096)	16777216(16781312)	4096
FC3	4096				1000	4096*1000 (+ 1000)	4096000(4097000)	1000

总共参数量为：138,344,128(138,375,544)

占用显存：

1. 参数占用显存：138344128×4/1024/1024=527.74MB～528MB（共有138344128个参数，每个都是float32类型的，即一个占用32位=4bytes）
2. 输出占用显存（和batch_size有关）：15237608×4/1024/1024=58.12MB/image