转自: 模型显存占用及其计算量_你吃过卤汁牛肉吗的博客-CSDN博客_请问有办法提前知道一个模型的所占显存吗
23.1.10更新: 节省显存王炸技巧:减少中间变量,尽量写成函数
其他节省显存的小技巧请参照:
把显存用在刀刃上!17 种 pytorch 节约显存技巧_听 风、的博客-CSDN博客_降低显存占用
Pytorch清空显存缓冲区(torch.cuda.empty_cache)_hxxjxw的博客-CSDN博客_pytorch清除gpu占用的内存
PyTorch 轻松节省显存的小技巧_一加一h的博客-CSDN博客_pytorch 清理显存
1. 显存的占用
当在GPU上跑一个模型时,显存的占用主要有两部分:
模型的输出(特征图、特征图的梯度)、模型的参数(权重矩阵、偏置值、梯度)
1. 模型参数的显存占用:(例如:卷积核的参数、BN层、全连接层的参数等(池化层没有参数))
2. 如果是在训练阶段,需要反向传播更新参数值,所以每个参数都需要存储梯度。所以模型参数的显存占用,与采用的优化器有关。
1)如果使用SGD,需要2倍模型参数的显存占用,模型参数+模型参数的梯度;
2)如果使用SGD+Momentum,需要3倍的模型参数显存占用,模型参数+梯度+动量;
3)如果使用Adam,需要4倍的显存占用,模型参数+梯度+动量+二阶动量。
3. 训练阶段,会保存模型每一层输出的特征图。(因为反向传播中需要对中间层的特征图求导,所以中间层的输出特征图不会被释放)
4. 训练阶段,会保存对于每一层输出特征图的导数。(因为反向传播中链式求导,公式中有),但是由于特征图不像模型参数一样需要优化更新,所以在训练阶段特征图需要2倍显存(特征图+特征图梯度)。
反向传播公式(以第2层为例):
,其中L为损失函数,为第2层输出的状态值,为第二层输出的激活值
显存的占用主要是模型本身的参数占用,还与每一层输出的特征图大小有关(这一部分与batch_size呈线性正相关)
2. 模型参数的显存占用
(1)每层卷积核的参数量计算公式:
(2)每层全连接的参数量计算公式:
(3)每层BN的参数量计算公式:
模型参数的显存占用与batchsize无关,但是在训练阶段与优化器有关
所以在预测阶段要使用
-
with torch.no_grad():
-
inputs = None
-
output = model(inputs)
从而减少显存的使用。
3. 每层特征图的显存占用
每层output特征图的显存占用为:,与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使得计算,则需要考虑前后两层,也即前一层的特征图通过计算得到当前层的输出,这其中所需要占用的计算量。
卷积层FLOPs计算
池化层FLOPs计算
注:上面的CHW都是指输出的特征图的值,根据输出值进行计算比较容易。
BN层FLOPs计算
全连接层FLOPs计算
减少卷积层计算量的方法:
1. 卷积核的分解
如:(1)用两个3*3的卷积堆叠代替一个7*7的卷积(达到相同的感受野)、
(2)用1*3和3*1的卷积代替3*3的卷积
2. 深度可分离卷积:先进行Deepwise卷积操作((H*W*1)* 卷积),再进行Pointwise卷积操作( (1*1*)* 卷积)
实例分析:
显存占用并不是严格的与batch size成正比,还与模型自身的参数和延伸出来的数据相关(每层特征图size、优化器等)
并不是batch size越大速度越快。在充分利用计算资源的时候,加大batch size在速度上的提升很有限。---也就是显存足够,但是GPU的计算能力达到上限。
比如batch size=8时,模型的FLOPs等于GPU的峰值FLOPS,并且假设GPU能够维持其峰值FLOPS,那么刚好1秒跑一个batch size;若此时令batch size=16,那么需要2秒跑一个batch size,并不能提升训练速度。