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

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，这些是神经网络的权严重小，可以看到第一层卷积是3x3，而输入图像的通道是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的时分产生的额外的中间参数

优化器在优化时产生的额外的模型参数

但其实，我们占用的显存空间为什么比我们实际计算的还要大，缘由大约是由于深度学习框架一些额外的开支吧，不过假设经过下面公式，实际计算出来的显存和实践不会差太多的。

如何优化

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

inplace 

撩我吧

假设你与我情投意合于此，老潘很情愿与你交流；

假设你喜欢老潘的内容，欢迎关注和支持。

假设你喜欢我的文章，希望点赞:+1: 收藏 :file_folder: 评论 :speech_balloon: 三连一下~

【编辑引荐】

Linux系统C言语编程基础视频课程

VBScript脚本言语编程与自动化运维操作学习篇

关于 Java 注解（annotation）编程

2021年值得学习的5种编程言语

曾扬言消灭人类的机器人竟沦为商演工具，人工智能终究还是太烧钱

(责任编辑：admin)