前言

亲，显存炸了，你的显卡快冒烟了！

torch.FatalError: cuda runtime error (2) : out of memory at /opt/conda/conda-bld/pytorch_1524590031827/work/aten/src/THC/generic/THCStorage.cu:58

想必这是所有炼丹师们最不想看到的谬误，没有之一。

OUT OF MEMORY，显然是显存装不下你那么多的模型权重还有两头变量，而后程序奔溃了。怎么办，其实方法有很多，及时清空两头变量，优化代码，缩小 batch，等等等等，都可能缩小显存溢出的危险。

然而这篇要说的是下面这所有优化操作的根底，如何去计算咱们所应用的显存。学会如何计算出来咱们设计的模型以及两头变量所占显存的大小，想必晓得了这一点，咱们对本人显存也就会得心应手了。

如何计算

首先咱们应该理解一下根本的数据量信息：

• 1 G = 1000 MB
• 1 M = 1000 KB
• 1 K = 1000 Byte
• 1 B = 8 bit

好，必定有人会问为什么是 1000 而不是 1024，这里不过多探讨，只能说两种说法都是正确的，只是利用场景略有不同。这里对立依照下面的规范进行计算。

而后咱们说一下咱们平时应用的向量所占的空间大小，以 Pytorch 官网的数据格式为例（所有的深度学习框架数据格式都遵循同一个规范）：

咱们只须要看右边的信息，在平时的训练中，咱们常常应用的个别是这两种类型：

• float32 单精度浮点型
• int32 整型

个别一个 8 -bit 的整型变量所占的空间为 1B 也就是 8bit。而 32 位的 float 则占4B 也就是32bit。而双精度浮点型 double 和长整型 long 在平时的训练中咱们个别不会应用。

ps：生产级显卡对单精度计算有优化，服务器级别显卡对双精度计算有优化。

也就是说，假如有一幅 RGB 三通道真彩色图片，长宽别离为 500 x 500，数据类型为单精度浮点型，那么这张图所占的显存的大小为：500 x 500 x 3 x 4B = 3M。

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

不多是吧，没关系，好戏才刚刚开始。

显存去哪儿了

看起来一张图片 (3x256x256) 和卷积层 (256x100x100) 所占的空间并不大，那为什么咱们的显存仍旧还是用的比拟多，起因很简略，占用显存比拟多空间的并不是咱们输出图像，而是神经网络中的两头变量以及应用 optimizer 算法时产生的巨量的两头参数。

咱们首先来简略计算一下 Vgg16 这个 net 须要占用的显存：

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

• 模型本身的参数(params)
• 模型计算产生的两头变量(memory)

图片来自 cs231n，这是一个典型的 sequential-net，自上而下很顺畅，咱们能够看到咱们输出的是一张 224x224x3 的三通道图像，能够看到一张图像只占用150x4k，但下面是150k，这是因为这里在计算的时候默认的数据格式是 8 -bit 而不是 32-bit，所以最初的后果要乘上一个 4。

咱们能够看到，右边的 memory 值代表：图像输出进去，图片以及所产生的两头卷积层所占的空间。咱们都晓得，这些不拘一格的深层卷积层也就是深度神经网络进行“思考”的过程：

图片从 3 通道变为 64 –> 128 –> 256 –> 512 …. 这些都是卷积层，而咱们的显存也次要是他们占用了。

还有下面左边的 params，这些是神经网络的权重大小，能够看到第一层卷积是 3 ×3，而输出图像的通道是 3，输入通道是 64，所以很显然，第一个卷积层权重所占的空间是 (3 x 3 x 3) x 64。

另外还有一个须要留神的是两头变量在 backward 的时候会翻倍！

举个例子，上面是一个计算图，输出x，通过两头后果z，而后失去最终变量L：

咱们在 backward 的时候须要保留下来的两头值。输入是 L，而后输出x，咱们在 backward 的时候要求L 对x的梯度，这个时候就须要在计算链 L 和x两头的z：

dz/dx这个两头值当然要保留下来以用于计算，所以粗略预计，backward的时候两头变量的占用了是 forward 的两倍！

优化器和动量

要留神，优化器也会占用咱们的显存！

为什么, 看这个式子:

上式是典型的 SGD 随机降落法的总体公式，权重 W 在进行更新的时候，会产生保留两头变量，也就是在优化的时候，模型中的 params 参数所占用的显存量会翻倍。

当然这只是 SGD 优化器，其余简单的优化器如果在计算时须要的两头变量多的时候，就会占用更多的内存。

模型中哪些层会占用显存

有参数的层即会占用显存的层。咱们个别的卷积层都会占用显存，而咱们常常应用的激活层 Relu 没有参数就不会占用了。

占用显存的层个别是：

• 卷积层，通常的 conv2d
• 全连贯层，也就是 Linear 层
• BatchNorm 层
• Embedding 层

而不占用显存的则是：

• 方才说到的激活层 Relu 等
• 池化层
• Dropout 层

具体计算形式：

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

额定的显存

总结一下，咱们在总体的训练中，占用显存大略分以下几类：

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

但其实，咱们占用的显存空间为什么比咱们实践计算的还要大，起因大略是因为深度学习框架一些额定的开销吧，不过如果通过下面公式，实践计算出来的显存和理论不会差太多的。

如何优化

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

• 缩小输出图像的尺寸
• 缩小 batch，缩小每次的输出图像数量
• 多应用下采样，池化层
• 一些神经网络层能够进行小优化，利用 relu 层中设置inplace
• 购买显存更大的显卡
• 从深度学习框架下面进行优化

撩我吧

• 如果你与我气味相投于此，老潘很违心与你交换；
• 如果你喜爱老潘的内容，欢送关注和反对。
• 如果你喜爱我的文章，心愿点赞???? 珍藏 ???? 评论 ???? 三连一下~

想晓得老潘是如何学习踩坑的，想与我交换问题~ 请关注公众号「oldpan 博客」。
老潘也会整顿一些本人的私藏，心愿能帮忙到大家，点击神秘传送门获取。