显存不足?PyTorch 显存使用分析与优化

面对动辄几百万几千万参数量的模型, GPU那连常规 U盘都比不过的显存, 真的是杯水车薪。相信大家在日常模型训练过程中,或多或少的总会遇见:

torch.FatalError: cuda runtime error (2) : out of memory 

那么, 可怜的显存到底在哪里消耗掉了?

显存去哪了?

首先我们看一下 PyTorch 官方的数据格式(由于文本只讨论显存,故将 CPU tensor 去掉了):

其中,大家日常训练用到的一般是torch.cuda.FloatTensor和torch.cuda.IntTensor, 根据上表,它们分别是32位浮点数和32位整数,都占4个 Byte.

也就是说,常见的一张224x224大小的图片,占用 个 Byte,约为588K。

根据常用的张量组织形式(bathc_size, channel, height, width), 那么一个大小时(128,3, 224, 224) 的张量约占用73.5M。

让我们来看一下常见的 VGG-16的内存占用:

上述图片来自 CS231n 的课件

(http://cs231n.stanford.edu/slides/2017/cs231n_2017_lecture9.pdf)

是 VGG-16 在输入单张图片,且在不计算 biases的情况下计算的内存和参数量占用。

红色字体代表:从输入图像开始,每一个的卷积池化,全连接所产升的内存消耗。

蓝色字体代表:每一步计算,所需要记住的参数的个数

注意,这只是前向 forward 时候的内存消耗,计算反向传播的时候,上述消耗要至少再乘以2,因为链式法则要记很多中间结果。

最终的结果有点沉重,单张图片96M,反向传播时单张图片192M。这时候,我们稍微升高一下 batch_size, 比如batch_size = 128, 结果可想而知

处理模型设计的固有内存占用,还有很多模块也占用非常多的显存。

比如:优化器

比如常用的随机梯度下降SGD,它的计算公式是:

显然,程序是要记忆

的,这一点在上面讲过。

但是,如果要使用动量(Momentum-SGD) , 程序还要额外记忆动量, 占用显存 x3:

而更常用的,如果我们使用 Adam 优化器,占用显存 x4。

总的来说,显存占用分四个部分:

  • 模型参数
  • 模型计算中间结果
  • 反向传播中间结果
  • 优化器额外参数
计算模型显存占用

比如我们有一个如下的 PyTorch 模型:

Sequential(
 (conv_1): Conv2d(3, 64, kernel_size=(3, 3), stride=(1, 1), 
padding=(1, 1))
 (relu_1): ReLU(inplace)
 (conv_2): Conv2d(64, 64, kernel_size=(3, 3), stride=(1, 1), 
padding=(1, 1))
 (relu_2): ReLU(inplace)
 (pool_2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, 
ceil_mode=False)
 (conv_3): Conv2d(64, 128, kernel_size=(3, 3), stride=(1, 1), 
padding=(1, 1))
)

我们可以使用model.parameters() 将模型的参数拿出来,然后统计一下个数:

import numpy asnp
# 模型内参数个数
para =sum([np.prod(list(p.size())) for p in model.parameters()])
# float32 占 4个 Byte
type_size =4
model_size =para * type_size /1024/1024
print(model_size)
# 约为0.43M

显然,这仅仅是第一部分,模型参数

模型计算的中间结果可以通过传入一个简单输入,然后提取中间结果来统计:

input_ =input.clone()  
# 取消计算梯度,我们现在只计算模型中间结果的大小
input_.requires_grad_(requires_grad=False)

mods =list(model.modules())
out_sizes =[]

fori inrange(1, len(mods)):
    m =mods[i]
    ifisinstance(m,nn.ReLU):  
        ifm.inplace:
           continue
    out =m(input_)
   out_sizes.append(np.array(out.size()))
   input_ = out

total_nums =0
fori inrange(len(out_sizes)):
    s =out_sizes[i]
    nums =np.prod(np.array(s))
   total_nums += nums
    

cost_size =total_nums * type_size /1000/1000
print(cost_size)
# 约为320.5M

正向计算结果约为320.5M,反向x2,约为641M

如何优化

出了模型层面的优化,其他的显存优化主要有以下几点:

  • 减少batch,减少每次的输入图像数量
  • 多使用下采样,池化层等
  • 一些神经网络层可以进行小优化,利用relu层中设置inplace

PyTorch 0.4推出了一个新功能,能够在一定程度上解决显存不足的问题,测量主要是拿时间换空间。它的实现方式是将一个计算过程分成很多份,我们就可以先计算一部分,保存后一部分需要的中间结果,然后再计算后一部分。

也就是说,新的checkpoint允许我们只存储反向传播所需要的部分内容。如果当中缺少一个输出(为了节省内存而导致的),checkpoint将会从最近的检查点重新计算中间输出,以便减少显存使用:

比如,我们有一个1000层的全连接:

input=torch.rand(1, 10)
# 1000 个 10 x 10 的全连接层
layers = [nn.Linear(10, 10) for _ inrange(1000)]
model = nn.Sequential(*layers)
output = model(input)

这个模型如果正常运行将占用海量显存,因为计算中会产生很多的中间变量。现在我们可以使用checkpoint节省资源占用。

input = torch.rand(1, 10, requires_grad=True)
layers = [nn.Linear(10, 10) for _ in range(1000)]

# 计算第一部分
def run_a(*args):
    x = args[0]
    for layer in layers[:500]:
        x = layer(x)
    return x

# 计算第二部分
def run_b(*args):
    x = args[0]
    for layer in layers[500:-1]:
        x = layer(x)
    return x

from torch.utils.checkpoint import checkpoint
result_a = checkpoint(run_a, input)
result_b = checkpoint(run_b, result_a)
total_result = layers[-1](result_b)
output = total_result.sum()
output.backward() 

-END-

原文发布于微信公众号 - 专知(Quan_Zhuanzhi)

原文发表时间:2018-06-21

本文参与腾讯云自媒体分享计划,欢迎正在阅读的你也加入,一起分享。

发表于

我来说两句

0 条评论
登录 后参与评论

相关文章

来自专栏目标检测和深度学习

如何从零开发一个复杂深度学习模型

深度学习框架中涉及很多参数,如果一些基本的参数如果不了解,那么你去看任何一个深度学习框架是都会觉得很困难,下面介绍几个新手常问的几个参数。 batch 深度学习...

5597

用Python的长短期记忆神经网络进行时间序列预测

长短期记忆递归神经网络具有学习长的观察序列的潜力。

3.1K9
来自专栏机器学习算法工程师

超快速!10分钟入门Keras指南

作者:李中粱 小编:赵一帆 1 Keras框架介绍 在用了一段时间的Keras后感觉真的很爽,所以特意祭出此文与我们公众号的粉丝分享。 Keras是一个非常方...

5688
来自专栏木东居士的专栏

漫谈机器学习之小知识点总结

2014
来自专栏大数据挖掘DT机器学习

如何用TensorFlow和TF-Slim实现图像标注、分类与分割

本文github源码地址: 在公众号 datadw 里 回复 图像 即可获取。 笔者将和大家分享一个结合了TensorFlow和slim库的小应用,来实现...

6414
来自专栏编程

关于反向传播在Python中应用的入门教程

我来这里的目的是为了测试我对于Karpathy的博客《骇客的神经网络指导》以及Python的理解,也是为了掌握最近精读的Derek Banas的文章《令人惊奇的...

2127
来自专栏素质云笔记

keras系列︱图像多分类训练与利用bottleneck features进行微调(三)

不得不说,这深度学习框架更新太快了尤其到了Keras2.0版本,快到Keras中文版好多都是错的,快到官方文档也有旧的没更新,前路坑太多。 到发文为...

1.6K8
来自专栏AI深度学习求索

如何计算显存的占用,常常遇到out of memory?

最近一次组会上,师兄点评一篇文章显存占用过多,突然发现还不知道如何具体的计算显存,只好去学习一下。

7911
来自专栏人工智能LeadAI

keras学习笔记-黑白照片自动着色的神经网络-Beta版

Alpha版本不能很好地给未经训练的图像着色。接下来,我们将在Beta版本中做到这一点——将上面的将神经网络泛化。 以下是使用Beta版本对测试图像着色的结果。...

3666
来自专栏人工智能LeadAI

简易的深度学习框架Keras代码解析与应用

总体来讲keras这个深度学习框架真的很“简易”,它体现在可参考的文档写的比较详细,不像caffe,装完以后都得靠技术博客,keras有它自己的官方文档(不过是...

7107

扫码关注云+社区

领取腾讯云代金券