机器学习中的GPU和显存

文章目录

1. 显存分析

1.1 存储指标

1.2 神经网络显存占用

1.2.1 参数的显存占用

1.2.2 梯度与动量的显存占用

1.2.3 输入输出的显存占用

1.3 节省显存的方法

2. 计算量分析

2.1 常用操作的计算量

2.2 AlexNet 分析

2.3 减少卷积层的计算量

2.4 常用模型 显存/计算复杂度/准确率

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

nvidia-smi的输出

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

• 显存占用

• GPU利用率

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

这里推荐一个好用的小工具：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. 显存分析

1.1 存储指标

K、M，G，T是以1024为底，而KB 、MB，GB，TB以1000为底。不过一般来说，在估算显存大小的时候，我们不需要严格的区分这二者。

在深度学习中会用到各种各样的数值类型，数值类型命名规范一般为TypeNum，比如Int64、Float32、Double64。

• Type：有Int，Float，Double等

• Num: 一般是 8，16，32，64，128，表示该类型所占据的比特数目

常用的数值类型如下图所示：

常用的数值类型

其中Float32 是在深度学习中最常用的数值类型，称为单精度浮点数，每一个单精度浮点数占用4Byte的显存。

举例来说：有一个1000x1000的 矩阵，float32，那么占用的显存差不多就是

2x3x256x256的四维数组（BxCxHxW）占用显存为：24M

1.2 神经网络显存占用

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

• 模型自身的参数

• 模型的输出

举例来说，对于如下图所示的一个全连接网络(不考虑偏置项b)

模型的输入输出和参数

模型的显存占用包括：

• 参数：二维数组 W

• 模型的输出：二维数组 Y

输入X可以看成是上一层的输出，因此把它的显存占用归于上一层。

这么看来显存占用就是W和Y两个数组？

并非如此！！！

下面细细分析。

1.2.1 参数的显存占用

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

有参数的层主要包括：

• 卷积

• 全连接

• 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()之后就会占用相应的显存，占用的显存大小基本与上述分析的显存差不多（会稍大一些，因为其它开销）。

1.2.2 梯度与动量的显存占用

举例来说， 优化器如果是SGD：

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

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

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

• 参数 W

• 梯度 dW（一般与参数一样）

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

1.2.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）

1.3 节省显存的方法

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

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

• 降低batch-size

• 下采样(NCHW -> (1/4)*NCHW)

• 减少全连接层（一般只留最后一层分类用的全连接层）

2. 计算量分析

计算量的定义，之前已经讲过了，计算量越大，操作越费时，运行神经网络花费的时间越多。

2.1 常用操作的计算量

常用的操作计算量如下：

• 全连接层：BxMxN , B是batch size，M是输入形状，N是输出形状。

卷积的计算量分析

• ReLU的计算量：BHWC

2.2 AlexNet 分析

AlexNet的分析如下图，左边是每一层的参数数目（不是显存占用），右边是消耗的计算资源

AlexNet分析

可以看出：

• 全连接层占据了绝大多数的参数

• 卷积层的计算量最大

2.3 减少卷积层的计算量

今年谷歌提出的MobileNet，利用了一种被称为DepthWise Convolution的技术，将神经网络运行速度提升许多，它的核心思想就是把一个卷积操作拆分成两个相对简单的操作的组合。如图所示, 左边是原始卷积操作，右边是两个特殊而又简单的卷积操作的组合（上面类似于池化的操作，但是有权重，下面类似于全连接操作）。

Depthwise Convolution

这种操作使得：

2.4 常用模型 显存/计算复杂度/准确率

去年一篇论文(http://link.zhihu.com/?target=https%3A//arxiv.org/abs/1605.07678)总结了当时常用模型的各项指标，横座标是计算复杂度（越往右越慢，越耗时），纵座标是准确率（越高越好），圆的面积是参数数量（不是显存占用）。左上角我画了一个红色小圆，那是最理想的模型的的特点：快，效果好，占用显存小。

常见模型计算量/显存/准确率