深度学习中GPU和显存分析

一 神经网络显存占用分析

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

• 模型自身的参数
• 模型的输出

1.1 参数的显存占用

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

（1）有参数的层主要包括：

• 卷积
• 全连接
• BatchNorm  BN层也是有参数的哦，
• Embedding层
• ... ...

（2）无参数的层：

• 多数的激活层(Sigmoid/ReLU)
• 池化层
• Dropout
• ... ...

（3）有参数的层的明细

更具体的来说，模型的参数数目(这里均不考虑偏置项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()之后就会占用相应的显存，占用的显存大小基本与上述分析的显存差不多（会稍大一些，因为其它开销）。

1.2 梯度与动量的显存占用

举例来说，

（1）优化器如果是SGD：除了保存W之外还要保存对应的梯度 ，因此显存占用等于参数占用的显存x2,

（2）如果是带Momentum-SGD，这时候还需要保存动量， 因此显存x3

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

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

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

1.3 输入输出的显存占用

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

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）

 

二、节省显存的方法

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

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

• 下采样(NCHW -> (1/4)*NCHW)
• 减少全连接层（一般只留最后一层分类用的全连接层）
• 减少输入图像的尺寸
• 减少batch，减少每次的输入图像数量
• 多使用下采样，池化层
• 一些神经网络层可以进行小优化，利用relu层中设置inplace（这个需要注意！）
• 购买显存更大的显卡
• 从深度学习框架上面进行优化

三、分析性能常用的工具

（1）nvidia-smi 命令

nvidia-smi是Nvidia显卡命令行管理套件，基于NVML库，旨在管理和监控Nvidia GPU设备。

nvidia-smi的输出

这是nvidia-smi命令的输出，其中最重要的两个指标：

• 显存占用
• GPU利用率

显存占用和GPU利用率是两个不一样的东西，显卡是由GPU计算单元和显存等组成的，显存和GPU的关系有点类似于内存和CPU的关系。

（2）gpustat 工具

gpustat,直接

pip install gpustat

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

watch --color -n1 gpustat -cpu 

gpustat 输出

显存可以看成是空间，类似于内存。

• 显存用于存放模型，数据
• 显存越大，所能运行的网络也就越大

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

1*2+3                  1 flop
1*2 + 3*4 + 4*5        3 flop 

（3）Linux下的 htop 工具

在说之前先推荐一个实时监控内存显存使用的小工具：

sudo apt-get install htop

监控内存（-d为更新频率，下为每0.1s更新一次）：

htop -d=0.1

监控显存（-n为更新频率，下为每0.1s更新一次）：

watch -n 0.1 nvidia-smi

（4）pynvml

这是Nvidia的Python环境库和Python的垃圾回收工具

可以实时地打印我们使用的显存以及哪些Tensor使用了我们的显存

https://github.com/Oldpan/Pytorch-Memory-Utils

（5）Pytorch-Memory-Utils工具

GitHub地址为：https://github.com/Oldpan/Pytorch-Memory-Utils

 

 

四、总结建议

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

尤其是batch-size，假定GPU处理单元已经充分利用的情况下：

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

4.1 关于显卡选购

当前市面上常用的显卡指标如下：