转自: https://blog.csdn.net/u012370185/article/details/95234125

1. 显存的占用

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

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

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

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

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

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

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

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

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

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

$\frac{\partial L}{\partial W_2}=\frac{\partial L}{\partial a_2} \cdot \frac{\partial a_2}{\partial z_2} \cdot \frac{\partial z_2}{\partial W_2} \\\\=-(y-a_2)^2f'(a_2)a_1$ ，其中L为损失函数， $z_2$ 为第2层输出的状态值， $a_2$ 为第二层输出的激活值

显存的占用主要是模型本身的参数占用，还与每一层输出的特征图大小有关（这一部分与batch_size呈线性正相关）

2. 模型参数的显存占用

(1)每层卷积核的参数量计算公式： $C_{out} \times(C_{in}\times K^2+1)$

(2)每层全连接的参数量计算公式： $FC_{in}\times FC_{out}$

(3)每层BN的参数量计算公式： $2\times C_{in}$

模型参数的显存占用与batchsize无关，但是在训练阶段与优化器有关

所以在预测阶段要使用

with torch.no_grad():
inputs = None
output = model(inputs)

从而减少显存的使用。

3. 每层特征图的显存占用

每层output特征图的显存占用为： $B\times C_{out}\times H\times W$ ，与batch size的大小成正比

训练阶段，需要保存特征图的梯度，所以需要2倍的显存占用

神经网络需要保存每一层的输入（前一层输出的激活值）输出（状态值），但是对于某些特殊情况，我们无需保存输入。比如采用ReLU作为激活函数时，使用

nn.ReLU(inplace=True)

能够将激活函数ReLU的输出值直接覆盖保存于模型的输出之中，节省了显存。感兴趣的读者可以思考⼀一下，这时候是如何反向传播的（提示：y=relu(x) ->dx = dy.copy();dx[y<=0]=0 ）`

4. 节省显存的几种方法

占用显存较多的部分是卷积层输出的特征图，模型参数的显存占用较少，其中全连接层占据了模型参数的大部分

所以降低模型显存一般有如下几种方法：

1. 去掉全连接层，改为全局平均池化

2. 降低batch size

3. 池化，减小特征图size

4. pytorch中的inplace操作，torch.no_grad()

5. FLOPs计算

FLOPS：全大写，是floating point operations per second的缩写，意指每秒浮点运算次数，理解为计算速度。是一个衡量硬件性能的指标。

FLOPs：注意s小写，是floating point operations的缩写（s表复数），意指浮点运算数，理解为计算量。可以用来衡量算法/模型的复杂度。

模型的FLOPs越大，那么执行一遍前向传播或者反向传播的时间也就越长，从而导致漫长的训练时间。

前面的显存占用，只需要考虑一层便足够了，比较好计算。而模型FLOP使得计算，则需要考虑前后两层，也即前一层的特征图通过计算得到当前层的输出，这其中所需要占用的计算量。