究竟什么占用了你的显存容量

电玩女神 2022-02-17 17:57 875阅读 0赞

## 1、先来看看平常使用的向量所占用空间的大小 ##

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xxejc5MDE5MjU5Mw_size_16_color_FFFFFF_t_70][] pytorch官方发布的向量占用

举例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)

![watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xxejc5MDE5MjU5Mw_size_16_color_FFFFFF_t_70 1][]

图片来自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`在进行更新的时候，会产生保存中间变量![《浅谈深度学习:如何计算模型以及中间变量的显存占用大小》][TIM_E6_88_AA_E5_9B_BE20180607104328.jpg]，也就是在优化的时候，模型中的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`
 *  购买显存更大的显卡
 *  从深度学习框架上面进行优化

[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xxejc5MDE5MjU5Mw_size_16_color_FFFFFF_t_70]: /images/20220217/6e8740f316014fce830cd9bc04757d81.png
[watermark_type_ZmFuZ3poZW5naGVpdGk_shadow_10_text_aHR0cHM6Ly9ibG9nLmNzZG4ubmV0L2xxejc5MDE5MjU5Mw_size_16_color_FFFFFF_t_70 1]: /images/20220217/6354d76f5eb24705bd654dfa5e7d5f7b.png
[TIM_E6_88_AA_E5_9B_BE20180607104328.jpg]: /images/20220217/26c02ce9962243fe9b113af36a39ef6f.png