模型显存占用及其计算量

1. 显存的占用

当在GPU上跑一个模型时，显存的占用主要有两部分：

模型的输出（特征图、特征图的梯度）、模型的参数（权重矩阵、偏置值、梯度）

1. 模型参数的显存占用：（例如：卷积核的参数、BN层、全连接层的参数等（池化层没有参数）） 

2. 如果是在训练阶段，需要反向传播更新参数值，所以每个参数都需要存储梯度。所以模型参数的显存占用，与采用的优化器有关。

         1）如果使用SGD，需要2倍模型参数的显存占用，模型参数+模型参数的梯度；

         2）如果使用SGD+Momentum，需要3倍的模型参数显存占用，模型参数+梯度+动量；

         3）如果使用Adam，需要4倍的显存占用，模型参数+梯度+动量+二阶动量。

3. 训练阶段，会保存模型每一层输出的特征图。（因为反向传播中需要对中间层的特征图求导，所以中间层的输出特征图不会被释放）

4. 训练阶段，会保存对于每一层输出特征图的导数。（因为反向传播中链式求导，公式中有），但是由于特征图不像模型参数一样需要优化更新，所以在训练阶段特征图需要2倍显存（特征图+特征图梯度）。

反向传播公式（以第2层为例）：