1、先来看看平常使用的向量所占用空间的大小
![](https://i-blog.csdnimg.cn/blog_migrate/fb57c030081486f1c1f66c1325ee3455.png)
举例32bit float tensor,占用4个字节,所以对于一张RGB三通道的图像来说,如果长宽分别为500 x 500,数据类型为单精度浮点型,那么这张图所占的显存的大小为:500 x 500 x 3 x 4B = 3M。一个(256,3,100,100)-(N,C,H,W)的FloatTensor所占的空间为256 x 3 x 100 x 100 x 4B = 31M。
2、显存究竟去了哪里?
占用显存比较多空间的并不是我们输入图像,而是神经网络中的中间变量以及使用optimizer算法时产生的巨量的中间参数。
我们首先来简单计算一下Vgg16这个net需要占用的显存:
通常一个模型占用的显存也就是两部分:
- 模型自身的参数(params)
- 模型计算产生的中间变量(memory)
图片来自cs231n,这是一个典型的sequential-net,自上而下很顺畅,我们可以看到我们输入的是一张224x224x3的三通道图像,可以看到一张图像只占用150x4k,但上面标注的是150k,这是因为上图中在计算的时候默认的数据格式是8-bit而不是32-bit,所以最后的结果要乘上一个4。
图片从3通道变为64 –> 128 –> 256 –> 512 …. 这些都是卷积层,而我们的显存也主要是他们占用了。 还有上面右边的params,这些是神经网络的权重大小,可以看到第一层卷积是3×3,而输入图像的通道是3,输出通道是64,所以很显然,第一个卷积层权重所占的空间是 (3 x 3 x 3) x 64。
另外还有一个需要注意的是中间变量在backward的时候会翻倍!
优化器也会占用我们的显存!
SGD随机下降法的总体公式,权重W
在进行更新的时候,会产生保存中间变量,也就是在优化的时候,模型中的params参数所占用的显存量会翻倍。
当然这只是SGD优化器,其他复杂的优化器如果在计算时需要的中间变量多的时候,就会占用更多的内存。
深度学习中神经网络的显存占用,我们可以得到如下公式:
显存占用 = 模型显存占用 + batch_size × 每个样本的显存占用
可以看出显存不是和batch-size简单的成正比,尤其是模型自身比较复杂的情况下:比如全连接很大,Embedding层很大
另外需要注意:
- 输入(数据,图片)一般不需要计算梯度
- 神经网络的每一层输入输出都需要保存下来,用来反向传播,但是在某些特殊的情况下,我们可以不要保存输入。比如ReLU,在PyTorch中,使用
nn.ReLU(inplace = True)
能将激活函数ReLU的输出直接覆盖保存于模型的输入之中,节省不少显存。感兴趣的读者可以思考一下,这时候是如何反向传播的(提示:y=relu(x) -> dx = dy.copy();dx[y<=0]=0)
3、总结
(1)模型中那些层会占用显存
有参数的层即会占用显存的层。我们一般的卷积层都会占用显存,而我们经常使用的激活层Relu没有参数就不会占用了。
占用显存的层一般是:
- 卷积层,通常的conv2d
- 全连接层,也就是Linear层
- BatchNorm层
- Embedding层
而不占用显存的则是:
- 刚才说到的激活层Relu等
- 池化层
- Dropout层
具体计算方式:
- Conv2d(Cin, Cout, K): 参数数目:Cin × Cout × K × K
- Linear(M->N): 参数数目:M×N
- BatchNorm(N): 参数数目: 2N
- Embedding(N,W): 参数数目: N × W
(2)额外的显存
总结一下,我们在总体的训练中,占用显存大概分以下几类:
- 模型中的参数(卷积层或其他有参数的层)
- 模型在计算时产生的中间参数(也就是输入图像在计算时每一层产生的输入和输出)
- backward的时候产生的额外的中间参数
- 优化器在优化时产生的额外的模型参数
但其实,我们占用的显存空间为什么比我们理论计算的还要大,原因大概是因为深度学习框架一些额外的开销吧,不过如果通过上面公式,理论计算出来的显存和实际不会差太多的。
4、优化
- 减少输入图像的尺寸
- 减少batch,减少每次的输入图像数量
- 多使用下采样,池化层
- 一些神经网络层可以进行小优化,利用relu层中设置
inplace
- 购买显存更大的显卡
- 从深度学习框架上面进行优化