深度学习 占用gpu内存 使用率为0_深度解析MegEngine亚线性显存优化技术

e67f066a99ad0c8896e0fe008eff1740.png

作者 | 旷视研究院 编辑 | Linda 基于梯度检查点的亚线性显存优化方法 [1] 由于较高的计算 / 显存性价比受到关注。MegEngine 经过工程扩展和优化,发展出一套行之有效的加强版亚线性显存优化技术,既可在计算存储资源受限的条件下,轻松训练更深的模型,又可使用更大 batch size,进一步提升模型性能,稳定 batchwise 算子。使用 MegEngine 训练 ResNet18/ResNet50,显存占用分别最高降低 23%/40%;在更大的 Bert 模型上,降幅更是高达 75%,而额外的计算开销几乎不变。该技术已在 MegEngine 开源,欢迎大家上手使用:https://github.com/MegEngine

深度神经网络训练是一件复杂的事情,它体现为模型的时间复杂度和空间复杂度,分别对应着计算和内存;而训练时内存占用问题是漂浮在深度学习社区上空的一块乌云,如何拨云见日,最大降低神经网络训练的内存占用,是一个绕不开的课题。

GPU 显卡等硬件为深度学习提供了必需的算力,但硬件自身有限的存储,限制了可训练模型的尺寸,尤其是大型深度网络,由此诞生出一系列相关技术,比如亚线性显存优化、梯度累加、混合精度训练、分布式训练,进行 GPU 显存优化。

其中,亚线性显存优化方法 [1] 由于较高的计算 / 显存性价比备受关注;旷视基于此,经过工程扩展和优化,发展出加强版的 MegEngine 亚线性显存优化技术,轻松把大模型甚至超大模型装进显存,也可以毫无压力使用大 batch 训练模型。

这里将围绕着深度学习框架 MegEngine 亚线性显存优化技术的工程实现和实验数据,从技术背景、原理、使用、展望等多个方面进行首次深入解读。

背  景

在深度学习领域中,随着训练数据的增加,需要相应增加模型的尺寸和复杂度,进行模型「扩容」;而 ResNet [2] 等技术的出现在算法层面扫清了训练深度模型的障碍。不断增加的数据和持续创新的算法给深度学习框架带来了新挑战,能否在模型训练时有效利用有限的~~ 计算~~ 存储资源,尤其是减少 GPU 显存占用,是评估深度学习框架性能的重要指标。

在计算存储资源一定的情况下,深度学习框架有几种降低显存占用的常用方法,其示例如下:

  • 通过合适的梯度定义,让算子的梯度计算不再依赖于前向计算作为输入,从而 in-place 地完成算子的前向计算,比如 Sigmoid、Relu 等;   

  • 在生命周期没有重叠的算子之间共享显存;   

  • 通过额外的计算减少显存占用,比如利用梯度检查点重新计算中间结果的亚线性显存优化方法 [1];   

  • 通过额外的数据传输减少显存占用,比如把暂时不用的数据从 GPU 交换到 CPU,需要时再从 CPU 交换回来。

上述显存优化技术在 MegEngine 中皆有不同程度的实现,这里重点讨论基于梯度检查点的亚线性显存优化技术。

原   理

一个神经网络模型所占用的显存空间大体分为两个方面: 

1)模型本身的参数。

2)模型训练临时占用的空间,包括参数的梯度、特征图等。其中最大占比是 2)中以特征图形式存在的中间结果,比如,从示例 [1] 可知,根据实现的不同,从 70% 到 90% 以上的显存用来存储特征图。

这里的训练过程又可分为前向计算,反向计算和优化三个方面,其中前向计算的中间结果最占显存,还有反向计算的梯度。第 1)方面模型自身的参数内存占用最小。 

MegEngine 加强版亚线性显存优化技术借鉴了 [1] 的方法,尤其适用于计算存储资源受限的情况,比如一张英伟达 2080Ti,只有 11G 的显存;而更贵的 Tesla V100,最大显存也只有 32G。

d329db9cfb12a873911b909b3cfefc68.png   图 1:亚线性显存优化原理,其中 (b) 保存了 Relu 结果,实际中 Relu 结果可用 in-place 计算

图 1(a) 给出了卷积神经网络的基本单元,它由 Conv-BN-Relu 组成。可以看到,反向计算梯度的过程依赖于前向计算获取的中间结果,一个网络需要保存的中间结果与其大小成正比,即显存复杂度为 O(n)。

本质上,亚线性显存优化方法是以时间换空间,以计算换显存,如图 1(b) 所示,它的算法原理如下:

  • 选取神经网络中 k 个检查点,从而把网络分成 k 个 block,需要注意的是,初始输入也作为一个检查点;前向计算过程中只保存检查点处的中间结果;   

  • 反向计算梯度的过程中,首先从相应检查点出发,重新计算单个 block 需要的中间结果,然后计算 block 内部各个 block 的梯度;不同 block 的中间结果计算共享显存。这种方法有着明显的优点,即大幅降低了模型的空间复杂度,同时缺点是增加了额外的计算:   

  • 显存占用从 O(n) 变成 O(n/k)+ O(k),O(n/k) 代表计算单个节点需要的显存,O(k) 代表 k 个检查点需要的显存, 取 k=sqrt(n),O(n/k)+ O(k)~O(sqrt(n)),可以看到显存占用从线性变成了亚线性;   

  • 因为在反向梯度的计算过程中需要从检查点恢复中间结果,整体需要额外执行一次前向计算。

工   程

在 [1] 的基础上,MegEngine 结合自身实践,做了工程扩展和优化,把亚线性显存优化方法扩展至任意的计算图,并结合其它常见的显存优化方法,发展出一套行之有效的加强版亚线性显存优化技术。

亚线性优化方法采用简单的网格搜索(grid search)选择检查点,MegEngine 在此基础上增加遗传算法,采用边界移动、块合并、块分裂等策略,实现更细粒度的优化,进一步降低了显存占用。

如图 2 所示,采用型号为 2080Ti 的 GPU 训练 ResNet50,分别借助基准、亚线性、亚线性 + 遗传算法三种显存优化策略,对比了可使用的最大 batch size。仅使用亚线性优化,batch size 从 133 增至 211,是基准的 1.6x;而使用亚线性 + 遗传算法联合优化,batch size 进一步增至 262,较基准提升 2x。

ba29dacbe944f49b46850c2a45b18039.png

图 2:三种显存优化方法优化 batch size 的对比:ResNet50

通过选定同一模型、给定 batch size,可以更好地观察遗传算法优化显存占用的情况。如图 3 所示,随着迭代次数的增加,遗传算法逐渐收敛显存占用,并在第 5 次迭代之后达到一个较稳定的状态。

3c7550413a1230ca9e5232325193d0cf.png

图 3:遗传算法收敛示意图

此外,MegEngine 亚线性优化技术通过工程改良,不再局限于简单的链状结构和同质计算节点, 可用于任意的计算图,计算节点也可异质,从而拓展了技术的适用场景;并可配合上述显存优化方法,进一步降低模型的显存占用。

实  验

MegEngine 基于亚线性显存技术开展了相关实验,这里固定 batch size=64,在 ResNet18 和 ResNet50 两个模型上,考察模型训练时的显存占用和计算时间。   

如图 4 所示,相较于基准实现,使用 MegEngine 亚线性显存技术训练 ResNet18 时,显存占用降低 32%, 计算时间增加 24%;在较大的 ReNet50 上,显存占用降低 40%,计算时间增加 25%。同时经过理论分析可知,模型越大,亚线性显存优化的效果越明显,额外的计算时间则几乎不变。8735441f4c83c4fcfc6c84e2fbf81e8b.png

图 4:MegEngine 亚线性优化技术实验显存 / 时间对比:ReNet18/ReNet50

在更大模型 Bert 上实验数据表明,借助 MegEngine 亚线性显存技术,显存占用最高降低 75%,而计算时间仅增加 23%,这与理论分析相一致。有兴趣的同学可前往 MegEngine ModeHub 试手更多模型实验:https://megengine.org.cn/model-hub/。

使  用

MegEngine 官网提供了亚线性显存优化技术的使用文档。当你的 GPU 显存有限,苦于无法训练较深、较大的神经网络模型,或者无法使用大 batch 进一步提升深度神经网络的性能,抑或想要使 batchwise 算子更加稳定,那么,MegEngine 亚线性显存优化技术正是你需要的解决方案。

上手 MegEngine 亚线性优化技术非常便捷,无需手动设定梯度检查点,通过几个简单的参数,轻松控制遗传算法的搜索策略。具体使用时,在 MegEngine 静态图接口中调用 SublinearMemoryConfig 设置 trace 的参数 sublinear_memory_config,即可打开亚线性显存优化:

from megengine.jit import trace, SublinearMemoryConfig

config = SublinearMemoryConfig()

@trace(symbolic=True, sublinear_memory_config=config)
def train_func(data, label, *, net, optimizer):
...

MegEngine 在编译计算图和训练模型时,虽有少量的额外时间开销,但会显著缓解显存不足问题。下面以 ResNet50 为例,说明 MegEngine 可有效突破显存瓶颈,训练 batch size 从 100 最高增至 200:

import os
from multiprocessing import Process


def train_resnet_demo(batch_size, enable_sublinear, genetic_nr_iter=0):
import megengine as mge
import megengine.functional as F
import megengine.hub as hub
import megengine.optimizer as optim
from megengine.jit import trace, SublinearMemoryConfig
import numpy as np

print(
"Run with batch_size={}, enable_sublinear={}, genetic_nr_iter={}".format(
batch_size, enable_sublinear, genetic_nr_iter
)
)
# 使用 GPU 运行这个例子
assert mge.is_cuda_available(), "Please run with GPU"
try:
# 我们从 megengine hub 中加载一个 resnet50 模型。
resnet = hub.load("megengine/models", "resnet50")

optimizer = optim.SGD(resnet.parameters(), lr=0.1,)

config = None
if enable_sublinear:
config = SublinearMemoryConfig(genetic_nr_iter=genetic_nr_iter)

@trace(symbolic=True, sublinear_memory_config=config)
def train_func(data, label, *, net, optimizer):
pred = net(data)
loss = F.cross_entropy_with_softmax(pred, label)
optimizer.backward(loss)

resnet.train()
for i in range(10):
batch_data = np.random.randn(batch_size, 3, 224, 224).astype(np.float32)
batch_label = np.random.randint(1000, size=(batch_size,)).astype(np.int32)
optimizer.zero_grad()
train_func(batch_data, batch_label, net=resnet, optimizer=optimizer)
optimizer.step()
except:
print("Failed")
return

print("Sucess")


# 以下示例结果在 2080Ti GPU 运行得到,显存容量为 11 GB

# 不使用亚线性内存优化,允许的 batch_size 最大为 100 左右
p = Process(target=train_resnet_demo, args=(100, False))
p.start()
p.join()
# 报错显存不足
p = Process(target=train_resnet_demo, args=(200, False))
p.start()
p.join()

# 使用亚线性内存优化,允许的 batch_size 最大为 200 左右
p = Process(target=train_resnet_demo, args=(200, True, 20))
p.start()
p.join(
展  望

如上所述,MegEngine 的亚线性显存优化技术通过额外做一次前向计算,即可达到 O(sqrt(n)) 的空间复杂度。如果允许做更多次的前向计算,对整个网络递归地调用亚线性显存算法,有望在时间复杂度为 O(n log n) 的情况下,达到 O(log n) 的空间复杂度。

更进一步,MegEngine 还将探索亚线性显存优化技术与数据并行 / 模型并行、混合精度训练的组合使用问题,以期获得更佳的集成效果。最后,在 RNN 以及 GNN、Transformer 等其他类型网络上的使用问题,也是 MegEngine 未来的一个探索方向。

了解更多信息请查询:

  • MegEngine GitHub:https://github.com/MegEngine

  • MegEngine 官网:https://megengine.org.cn

  • MegEngine ModelHub:https://megengine.org.cn/model-hub/

参考文献

  1. Chen, T., Xu, B., Zhang, C., & Guestrin, C. (2016). Training deep nets with sublinear memory cost. arXiv preprint arXiv:1604.06174.  

  2. He, K., Zhang, X., Ren, S., & Sun, J. (2016). Deep residual learning for image recognition. In Proceedings of the IEEE conference on computer vision and pattern recognition (pp. 770-778).

2f53e982ca1d1e6fab2a3822a9d31336.gif

你也「在看」吗??

本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若转载,请注明出处:http://www.mzph.cn/news/297857.shtml

如若内容造成侵权/违法违规/事实不符,请联系多彩编程网进行投诉反馈email:809451989@qq.com,一经查实,立即删除!

相关文章

jq select操作全集

添加option $("#ID option").each(function(){if($(this).val()111){$(this).remove();}});移除option $("<option value111>UPS Ground</option>").appendTo($("#ID"));取得下拉选单的选取值 $("#testSelect option:selected&…

有趣的物理照片,让你瞬间爱上物理!

全世界只有3.14 % 的人关注了爆炸吧知识让你爱上科学的神奇物理现象图&#xff01;烧不坏的毛巾物理学好让人敬畏…不明觉厉啊&#xff01;重心求超越的吗&#xff1f;依旧是一个好玩的实验永...永动鸡&#xff1f;有种别划桨呀头朝下的飞机中&#xff0c;依旧可以正常倒水光的…

suse linux 如何修改主机名,suse如何修改主机名

在SLES8中如何修改主机名(hostname)Q: 在SLES8中如何修改主机名(hostname)A: SLES8中可以用yast network来修改主机名(hostname),也可以通过修改配置文件的方式并运行命令实现这个目的方法1. 通过yast2管理工具进行配置在console中运行yast2工具修改主机名(hostname)# yast2 ne…

SSH项目中遇到拦截器无法注入服务的问题

2019独角兽企业重金招聘Python工程师标准>>> 配置applicationContext.xml&#xff0c;给自定义拦截器增加了properity属性是一个biz&#xff0c;运行的时候一直报空指针异常&#xff0c;输出结果biz总是null&#xff0c;很是好奇&#xff0c;难不成拦截器无法被注入…

【转】定位oops的具体代码行

最近在学写linux驱动&#xff0c;按照宋宝华的《linux设备驱动开发详解》上的例子一步步来写。编译运行以后出现oops&#xff0c;不知道怎么调试&#xff0c;上网搜了一下&#xff0c;看到这篇文章&#xff0c;感觉不错转载到这里。 来自Linus Torvalds的讨论: https://groups.…

Dapr牵手.NET学习笔记:用docker-compose部署服务

上一篇聊到用两个物理机&#xff08;一个win,一个mac&#xff09;来部署dapr和服务 &#xff0c;实现order调用pay的负载均衡。本篇说一下在windows上的docker部署这三个服务&#xff0c;达到与上一篇的效果。三个服务的部署架构是这样的首先要把OrderSystem(服务端口80)项目do…

Java并发编程实战_真香!阿里P8耗时半年著作660页Java高并发与网络编程实战总结...

随着软件行业的飞速发展&#xff0c;互联网公司对开发者的技能要求也越来越高。而高并发、网络编程、微服务、海量数据的处理等技能&#xff0c;是每一个开发者进阶时的必学知识。为了帮助初级开发者快速掌握这些实用技术&#xff0c;本书以“理论&#xff0b;范例”的形式对各…

oracle如何往dg加盘_oracle 在物理机上添加磁盘操作

物理机上添加磁盘操作注意&#xff1a;1)物理机上添加磁盘操作&#xff0c;不涉及到start_udev的动作。2)磁盘分区的操作&#xff0c;需要谨慎进行&#xff0c;核准无误后再操作。(1)查看磁盘名称命名# su - grid$ sqlplus / as sysasmset linesize 180col name format a20col …

【译】在Asp.Net中操作PDF – iTextSharp - 使用表格

使用Asp.Net生成PDF最常用的元素应该是表格&#xff0c;表格可以帮助比如订单或者发票类型的文档更加格式化和美观。本篇文章并不会深入探讨表格&#xff0c;仅仅是提供一个使用iTextSharp生成表格的方法介绍&#xff0c;本文需要阅读我之前iTextSharp系列文章作为基础&#xf…

linux docker导入镜像,Docker镜像的导入和导出

相关阅读&#xff1a;场景描述&#xff1a;需要在客户现场快速部署应用&#xff0c;东西多&#xff0c;时间短场景分析&#xff1a;为了节省时间&#xff0c;使用docker进行快速部署&#xff0c;由于不是内部环境&#xff0c;无法使用内部私有库&#xff0c;于是构建镜像&#…

有些图,只要看错一眼就再也回不去了!

全世界只有3.14 % 的人关注了爆炸吧知识平时在网络上逛&#xff0c;你会发现有种图片&#xff0c;一旦看茬&#xff0c;就再也回不去了&#xff01;今天就来集体复习一下&#xff1a;首先是这张著名的图&#xff0c;很可爱的小盆友但如果告诉你&#xff1a;那两个大黑点是鼻孔&…

2015年网页设计最佳颜色搭配的9种选择

2019独角兽企业重金招聘Python工程师标准>>> 2015年网页设计最佳颜色搭配的9种选择 还在为你的网站选用哪种色调发愁&#xff0c;难易选择吗&#xff1f;本篇为你总结了2015年国外网友设计最佳颜色搭配的9个方案&#xff0c;供你参考... 详细解读 和小伙伴们一起来吐…

CryEngine

最近CryEngine3已经发布了&#xff08;实际上已经到3.1了&#xff09;&#xff0c;看到很多朋友都想要拿到最新的版本&#xff0c;在这里附上具体的申请方法。楼主看到了&#xff0c;麻烦看看能不能置个顶。 CryEngine对于教育机构是有免费授权的。实际上授权也十分宽松&#x…

EF Core的一个紧急bug,我这样修改

1背景今日在生产环境碰到如下错误ASP.NET MVC项目 Repository层中&#xff0c;Delete总是失败another entity of the same type already has the same primary key value具体错误提示&#xff1a;Attaching an entity of type ResearchManager.Models.BigTracker_UI.Product_Tr…

电脑home键在哪_如何灵活使用电脑键盘上的各个键

电脑上的键有很多但是你真的了解它们的用法么&#xff0c;今天笔者给大家分享一下电脑键盘上各个键的作用。区域一&#xff1a;共有13个按键&#xff0c;ESC&#xff0c;F1-F2。F1&#xff1a;帮助信息。F2&#xff1a;选定一个文件或文件夹。按下F2可以重命名。F3&#xff1a;…

svn强制注释 linux,svn强制要求提交注释

看了N多资料&#xff0c;不知道为什么我总是不成功。现在终于测试成功了&#xff0c;下面是实际操作过程~~使用bitnami一键安装了subversion&#xff0c;在使用中&#xff0c;希望开发人员提交时必须输入日志内容&#xff0c;可通过以下方法实现。首先转到相应的库的hook目录中…

jquery easyui datagrid getSelections用法

1.datagrid、 数据绑定 $(#tt).datagrid({ url: GetDataJosn, title: DataGrid, width: 800, height: 300, pageSize: 10, idField: productid, fitColumns: true,…

XCode5 真机调试及发布应用

一、XCODE 真机测试 Xcode5已经很智能&#xff0c;只需生成一个开发证书&#xff0c;安装后&#xff0c;插入设备会自动添加&#xff0c;注意&#xff0c;当Mac系统升级后&#xff0c;证书需要重新生成。证书生成步骤&#xff1a;1、生成 CertificateSigningRequest.certSignin…

身家4400亿美元的他,吃低于3美元的早餐,和2个老婆同居28年!

全世界只有3.14 % 的人关注了爆炸吧知识1930年8月30日&#xff0c;巴菲特出生在美国内布拉斯加州的奥马哈市。当时的美国还笼罩在经济大萧条的阴影中&#xff0c;还好父亲是位出色的股票掮客&#xff0c;这也让小巴菲特的童年过得顺顺利利。但小巴同学从小就不合群&#xff0c;…

怎么才能把项目做烂?!

上一篇聊聊《为什么被用户牵着鼻子走&#xff1f;刚参加工作做项目常说的一句话&#xff1a;没有困难&#xff0c;制造困难也得上。到不是真闲得蛋疼制造困难&#xff0c;而是表达解决困难的决心。2009年在某煤炭集团做了一个动态监测煤质煤量的实时系统&#xff0c;当时一般为…