深度模型GPU内存知识点

如果想要在单块GPU上训练一个大于VGG-16的模型，也许有几个解决内存限制问题的选择。

• 减少批量大小，但这会妨碍训练速度和精确度
• 在多GPU中分布模型
• 缩小模型
• 等待下一代GPU

什么占用内存？

可以根据功能性把GPU内存中的数据分为4个部分：

• 模型参数（权重）
• 特征映射
• 梯度映射
• 工作区

特征映射是正向过程中生成的中间结果，梯度映射是反向过程中生成的中间结果。工作区是cuDNN函数临时变量/矩阵的一个缓冲区。

一般而言，网络层数多，分配给特征映射的内存比例就越高，对于像VGG-16这样更大的模型来说，这一比例超过50%。

参数的显存占用

只有有参数的层，才会有显存占用，这部分的显存占用和输入无关，模型加载完之后就会占用

有参数的层主要包括：

• 卷积
• 全连接
• BatchNorm
• Embedding层
• ......

无参数的层

• 多数的激活层
• 池化层
• Dropout
• ......

模型的参数数目（不考虑偏置项b）：

• Linear(M->N)：参数数目 M*N
• Conv2d(Cin,Count,K):参数数目 CinCountK*K
• BatchNorm(N)：参数数目 2N
• Embedding(N,W):参数数目 N*W

参数占用显存 = 参数数目*n

• n=4 float32
• n=2 float16
• n=8 double64

梯度与动量的显存占用

举例来说，optimizer为SGD:

Wt+1 = Wt = αΔF(Wt)

显而易见，除了保存W之外，还要保存对应的梯度，因此显存占用等于参数占用的显存*2

如果是带Momentum-SGD,这个时候还要保存动量，显存*3

如果是Adam优化器，动量占用的显存更多，显存*4

总结以上，模型中与输入无关的显存占用包括：

• 参数W
• 梯度dW
• 优化器的动量（普通SGD没有动量，momentum-SGD动量与梯度一样，Adam优化器动量的数量是梯度的两倍）

输入输出的显存占用

这部分的显存主要看输出的feature map的形状。

image

比如卷积的输入输出满足一下关系：

image

根据以上可以计算每一层输出的Tensor的形状，然后就能计算相应的显存占用。

模型输出的显存占用，总结如下：

• 需要计算每一层的feature map的形状 （多维数组的形状）
• 需要保存输出对应的梯度用于反向传播（链式法则）
• 显存占用与batch size成正比
• 模型输出不需要存储相应的动量信息

深度学习中神经网络的显存占用，有以下公式：

显存占用 = 模型显存占用  + batch_size * 每个样本的显存占用

节省显存的方法

在深度学习中，一般占用显存最多的是卷基层的输出，模型参数占用的显存相对较少，而且不太好优化。

节省显存一般有如下方法：

• 降低batch-size
• 下采样(NCHW -> (1/4)*NCHW)
• 减少全连阶层 （一般只留最后一层分类用的全连阶层）

计算量分析

常用操作的计算量

• 全连阶层 BMN B是batch size,M是输入形状，N是输出形状
• 卷积的计算量 BHWCoutCinKK

image

• BatchNorm 计算量估算大概 BHWC*｛4，5，6｝
• 池化的计算量 BHWCK*K

image

• ReLU的计算量 BHWC

AlexNet分析

左边是每一层的参数数目，右边是消耗的计算量。

image

可以看出：

• 全连阶层占据了绝大多数的参数
• 卷基层的计算量最大

减少卷基层的计算量

MobileNet，利用DepthWise Convolution的技术，将神经网络运行速度提升很多，把一个卷积操作拆分成两个相对简单的操作的组合。左边是原始卷积操作，右边是两个特殊而又简单的卷积操作的组合（上面类似于池化的操作，但是有权重，下面类似于全连接操作）。

image

• 显存占用变多
• 计算量变少了许多，变成原来的 （1/Cout + 1/(k*k)）,一般为10-15%