深度学习（CV）显存占用分析

深度学习（CV）显存占用分析

一般来说，占用显存大的并不是输入图像，而是神经网络中的中间变量以及使用optimizer算法时产生的巨量的中间参数。参考https://oldpan.me/archives/how-to-calculate-gpu-memory

占用显存大概分以下几类：

• 模型中的参数(卷积层或其他有参数的层)
• 模型在计算时产生的中间参数(也就是输入图像在计算时每一层产生的输入和输出)
• backward的时候产生的额外的中间参数
• 优化器在优化时产生的额外的模型参数

模型参数+forward时的中间变量

1、假如输入的图像大小为 3 x 224 x 224 ，batch_size=16 ，数据用的是float32，每个数字占用4B，那么输入部分所占用的显存为16 x 3 x 224 x 224 x 4B =9.1875 MB

2、一个(256,3,100,100)-(N,C,H,W)的卷积核FloatTensor所占的空间为256 x 3 x 100 x 100 x 4B = 31M

通常一个模型占用的显存也就是两部分：

• 模型自身的参数(params)
• 模型计算产生的中间变量(memory)

以Vgg16为例：

（图中默认用1B空间来存储数据，实际应该是4B，所以应该多乘以4才对）

左边的memory值代表：图像输入进去，图片以及所产生的中间feature map所占的空间。

右边的params值代表：模型的参数量，比如一个Conv对应(3x3x3)x64的参数量，因为它是一个3x3的卷积核。卷积核的权重值也要放在显存上，所以也占用空间。而poolong层是没有参数的，只是定义了计算方式而已。

总结一下，模型中会占用显存的有：

占用显存：

• 卷积层，通常的conv2d
• 全连接层，也就是Linear层。（Embedding层就是以one hot为输入、中间层节点为字向量维数的全连接层。）
• BatchNorm层

不占用显存：

• Relu等激活函数
• 池化层
• Dropout层

具体计算方式：

• Conv2d(Cin, Cout, K): 参数数目：Cin × Cout × K × K
• Linear(M->N): 参数数目：M×N
• BatchNorm(N): 参数数目： 2N

backward时的中间变量

优化器做梯度下降时占用的显存

优化方法

优化除了算法层的优化，最基本的优化无非也就一下几点：

• 减少输入图像的尺寸
• 减少batch，减少每次的输入图像数量
• 多使用下采样，池化层