深入理解卷积与模型大小问题，解决显存不足

Posted 2021-03-18 yumoye

深入理解卷积与模型大小问题，解决显存不足

在训练自己的模型时常常出现显存不足等问题，这个时候我们常用的方法就是调参。一般常用的方法有以下几点：

模型压缩
网络参数调整， 比如减小训练图像大小，降低FC output个数，使用小的conv kernel size等。
深度学习框架调整
减层
但是对于既定网络，我们减小训练图像大小或者改变batchsize大小都会影响模型的性能，这点在目标检测和语义分割中反应比较明显。那么我们要做的是在不改变网络性能的情况下，尽量的压缩模型，空出足够的显存进行GPU加速。

GPU基础知识
我们利用gpustat工具包可以看到关于GPU的相关信息:

pip install gpustat

直接即可安装，gpustat基于nvidia-smi，可以提供更美观简洁的展示，结合watch命令，可以动态实时监控GPU的使用情况。

watch --color -n1 gpustat -cpu

 

GPU计算单元类似于CPU中的核，用来进行数值计算。衡量计算量的单位是flop： the number of floating-point multiplication-adds，浮点数先乘后加算一个flop。计算能力越强大，速度越快。衡量计算能力的单位是flops： 每秒能执行的flop数量

1Byte = 8 bit
1K = 1024 Byte
1M = 1024 K
1G = 1024 M
1T = 1024 G

10 K = 10*1024 Byte
卷积神经网络参数
在图像处理中，往往把图像表示为像素的向量，比如一个1000×1000的图像，可以表示为一个1000000的向量。如果隐含层数目与输入层一样，即也是1000000时，那么输入层到隐含层的参数数据为1000000×1000000=10^12，这样就太多了，基本没法训练。卷积神经网络通过CNN中的局部连接(Sparse Connectivity)和权值共享(Shared Weights)来减少参数。

神经网络模型占用的显存包括：

模型自身的参数
模型的输出
举例来说，对于如下图所示的一个全连接网络(不考虑偏置项b)

                                                          

模型的显存占用包括：

参数：二维数组 W
模型的输出： 二维数组 Y
其他
输入X可以看成是上一层的输出，因此把它的显存占用归于上一层。

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

有参数的层主要包括：

卷积
全连接
BatchNorm
Embedding层
... ...
无参数的层：

多数的激活层(Sigmoid/ReLU)
池化层
Dropout
... ...
更具体的来说，模型的参数数目(这里均不考虑偏置项b)为：

Linear(M->N): 参数数目：M×N
Conv2d(Cin, Cout, K): 参数数目：Cin × Cout × K × K
BatchNorm(N): 参数数目： 2N
Embedding(N,W): 参数数目： N × W
参数占用显存 = 参数数目×n

n = 4 ：float32

n = 2 : float16

n = 8 : double64

在PyTorch中，当你执行完model=MyGreatModel().cuda()之后就会占用相应的显存，占用的显存大小基本与上述分析的显存差不多（会稍大一些，因为其它开销）。

梯度与动量的显存占用
举例来说， 优化器如果是SGD：

                                                              

 

可以看出来，除了保存W之外还要保存对应的梯度，

                                                                            

因此显存占用等于参数占用的显存x2,

如果是带Momentum-SGD

                                                                

这时候还需要保存动量， 因此显存x3

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

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

参数 W
梯度 dW（一般与参数一样）
优化器的动量（普通SGD没有动量，momentum-SGD动量与梯度一样，Adam优化器动量的数量是梯度的两倍）
输入输出的显存占用
这部份的显存主要看输出的feature map 的形状。

                                              

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

                                      

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

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

需要计算每一层的feature map的形状（多维数组的形状）
需要保存输出对应的梯度用以反向传播（链式法则）
显存占用与 batch size 成正比
模型输出不需要存储相应的动量信息。
深度学习中神经网络的显存占用，我们可以得到如下公式：

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

可以看出显存不是和batch-size简单的成正比，尤其是模型自身比较复杂的情况下：比如全连接很大，Embedding层很大

另外需要注意：

输入（数据，图片）一般不需要计算梯度
神经网络的每一层输入输出都需要保存下来，用来反向传播，但是在某些特殊的情况下，我们可以不要保存输入。比如ReLU，在PyTorch中，使用nn.ReLU(inplace = True) 能将激活函数ReLU的输出直接覆盖保存于模型的输入之中，节省不少显存。感兴趣的读者可以思考一下，这时候是如何反向传播的（提示：y=relu(x) -> dx = dy.copy();dx[y<=0]=0）
在通过卷积获得了特征 (features) 之后，下一步我们希望利用这些特征去做分类。理论上讲，人们可以用所有提取得到的特征去训练分类器，例如 softmax 分类器，但这样做面临计算量的挑战。例如：对于一个 96X96 像素的图像，假设我们已经学习得到了400个定义在8X8输入上的特征，每一个特征和图像卷积都会得到一个 (96 − 8 + 1) × (96 − 8 + 1) = 7921 维的卷积特征，由于有 400 个特征，所以每个样例 (example) 都会得到一个 7921 × 400 = 3,168,400 维的卷积特征向量。学习一个拥有超过 3 百万特征输入的分类器十分不便，并且容易出现过拟合 (over-fitting)。
为了解决这个问题，首先回忆一下，我们之所以决定使用卷积后的特征是因为图像具有一种“静态性”的属性，这也就意味着在一个图像区域有用的特征极有可能在另一个区域同样适用。因此，为了描述大的图像，一个很自然的想法就是对不同位置的特征进行聚合统计，例如，人们可以计算图像一个区域上的某个特定特征的平均值 (或最大值)。这些概要统计特征不仅具有低得多的维度 (相比使用所有提取得到的特征)，同时还会改善结果(不容易过拟合)。这种聚合的操作就叫做池化 (pooling)，有时也称为平均池化或者最大池化 (取决于计算池化的方法)。

参数细节

Convolution层：

就是卷积层，是卷积神经网络（CNN）的核心层。
层类型：Convolution
　　  lr_mult: 学习率的系数，最终的学习率是这个数乘以solver.prototxt配置文件中的base_lr。如果有两个lr_mult, 则第一个表示权值的学习率，第二个表示偏置项的学习率。一般偏置项的学习率是权值学习率的两倍。
在后面的convolution_param中，我们可以设定卷积层的特有参数。
必须设置的参数：
  　　num_output: 卷积核（filter)的个数
  　　kernel_size: 卷积核的大小。如果卷积核的长和宽不等，需要用kernel_h和kernel_w分别设定
其它参数：
  　　stride: 卷积核的步长，默认为1。也可以用stride_h和stride_w来设置。
  　　pad: 扩充边缘，默认为0，不扩充。 扩充的时候是左右、上下对称的，比如卷积核的大小为5*5，那么pad设置为2，则四个边缘都扩充2个像素，即宽度和高度都扩充了4个像素,这样卷积运算之后的特征图就不会变小。也可以通过pad_h和pad_w来分别设定。
    　 weight_filler: 权值初始化。 默认为“constant",值全为0，很多时候我们用"xavier"算法来进行初始化，也可以设置为”gaussian"
    　 bias_filler: 偏置项的初始化。一般设置为"constant",值全为0。
   　  bias_term: 是否开启偏置项，默认为true, 开启
   　  group: 分组，默认为1组。如果大于1，我们限制卷积的连接操作在一个子集内。如果我们根据图像的通道来分组，那么第i个输出分组只能与第i个输入分组进行连接。
输入：n*c0*w0*h0
输出：n*c1*w1*h1
其中，c1就是参数中的num_output，生成的特征图个数
 w1=(w0+2*pad-kernel_size)/stride+1;
 h1=(h0+2*pad-kernel_size)/stride+1;
如果设置stride为1，前后两次卷积部分存在重叠。如果设置pad=(kernel_size-1)/2,则运算后，宽度和高度不变。

Pooling层
也叫池化层，为了减少运算量和数据维度而设置的一种层。
层类型：Pooling
必须设置的参数：
       kernel_size: 池化的核大小。也可以用kernel_h和kernel_w分别设定。
其它参数：
  　  pool: 池化方法，默认为MAX。目前可用的方法有MAX, AVE, 或STOCHASTIC
　　 pad: 和卷积层的pad的一样，进行边缘扩充。默认为0
　　 stride: 池化的步长，默认为1。一般我们设置为2，即不重叠。也可以用stride_h和stride_w来设置。

pooling层的运算方法基本是和卷积层是一样的。
输入：n*c*w0*h0
输出：n*c*w1*h1
和卷积层的区别就是其中的c保持不变
w1=(w0+2*pad-kernel_size)/stride+1;
h1=(h0+2*pad-kernel_size)/stride+1;
如果设置stride为2，前后两次卷积部分不重叠。100*100的特征图池化后，变成50*50.

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

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

降低batch-size
下采样(NCHW -> (1/4)*NCHW)
减少全连接层（一般只留最后一层分类用的全连接层）
减少卷积层的计算量
今年谷歌提出的MobileNet，利用了一种被称为DepthWise Convolution的技术，将神经网络运行速度提升许多，它的核心思想就是把一个卷积操作拆分成两个相对简单的操作的组合。如图所示, 左边是原始卷积操作，右边是两个特殊而又简单的卷积操作的组合（上面类似于池化的操作，但是有权重，下面类似于全连接操作）。

                                          Depthwise Convolution

这种操作使得：

显存占用变多(每一步的输出都要保存)
计算量变少了许多，变成原来的（ +）（一般为原来的10－15%）
常用模型 显存/计算复杂度/准确率
去年一篇论文(https://arxiv.org/abs/1605.07678)总结了当时常用模型的各项指标，横座标是计算复杂度（越往右越慢，越耗时），纵座标是准确率（越高越好），圆的面积是参数数量（不是显存占用），参数量越多，保存的模型文件越大。左上角我画了一个红色小圆，那是最理想的模型：快，准确率高，显存占用小。

                     

建议
时间更宝贵，尽可能使模型变快（减少flop）
显存占用不是和batch size简单成正比，模型自身的参数及其延伸出来的数据也要占据显存
batch size越大，速度未必越快。在你充分利用计算资源的时候，加大batch size在速度上的提升很有限
尤其是batch-size，假定GPU处理单元已经充分利用的情况下：

增大batch size能增大速度，但是很有限（主要是并行计算的优化）
增大batch size能减缓梯度震荡，需要更少的迭代优化次数，收敛的更快，但是每次迭代耗时更长。
增大batch size使得一个epoch所能进行的优化次数变少，收敛可能变慢，从而需要更多时间才能收敛（比如batch_size 变成全部样本数目）。