旷视开源「天元」,共创深度学习框架新纪元

 

3月25日,旷视召开Brain++ • 天元深度学习框架开源发布会,宣布「天元」开源。「天元」(MegEngine)始于2014年,是一个旷视自主研发的工业级深度学习框架,帮助旷视拿下27项全球AI竞赛冠军和MS COCO三连冠。旷视1400名研发人员使用它,开发公司的所有AI产品。

 

天元开源的代码有35万行,C++代码占大部分,其次是CUDA、Python;它采用Apache License2.0,方便大家使用,包括商业的场景。经过六年的工业验证之后, 天元针对这次开源做了特别的大型升级,以全新的面目在新一代人工智能开放平台OpenI和Gihtub上与大家见面。

 

天元的理念是「深度学习简单开发」,正好和旷视的基因「技术信仰价值务实」一一相映。天元之名取自围棋,以开源之名,向AlphaGO致敬,表明旷视秉持的技术信仰;同时天元寓意天地万物之元基,希望天元成为中国AI「新基建」的基石,彰显旷视坚守的价值务实。天元希望在技术与价值的双螺旋上升之中,与开发者一起,共创深度学习框架新纪元。

天元的整体架构

作为旷视全员使用的工业级深度学习框架,天元有着先进的整体架构,由上而下,它分为五个层次:1)「计算接口层」,2)「图表示层」,3)「优化与编译层」,4)「运行时管理层」,5)「计算内核层」。

 

天元架构示意图

最上面的「计算接口层」向外暴露了Python和C++接口,开发者借由这两种高级语言,使用天元编程,设计和研发原型系统,进行训练和推理。

 

第二层是「图表示」,它包含动态图和静态图的表示功能。第三层是「优化与编译」,它包括自动求导、图优化、图编译三个模块,是一个完整的、一体化的核心计算图,支撑着其上的计算接口、动态图、静态图的完整功能。

 

第四层是「运行时管理」,它包含计算调度和内存管理两个部分,前者把整个计算设备抽象为执行流,由调度器做合理调度;后者内置了管理动态内存、静态内存的高级优化模块,并实现了静态的亚线性内存优化器,极大提升内存管理效率。

 

最下面是支撑整个系统的「计算内核层」,内置高性能算子库,支持一系列常见计算设备,诸如X86、CUDA、ARM、ASIC等;同时还有高性能异构通信库,支持整个计算框架在分布式多节点上实现大规模训练。

天元的四大优势

天元是一个复杂的系统,它的目的是什么呢?正是为了解决旷视多年遇到的AI产品开发困境。我们主要讲四个困境,由此带出天元的四个优势——「训练推理一体」、「动静合一」、「兼容并包」、「灵活高效」。

 

通过它们,天元可以实现AI产品从实验原型到工业部署的小时级转化能力,支持大规模弹性训练,支持用更大数据、更大模型、更优算法生产更佳的AI产品,打造新一代AI生产力平台基础。

训练推理一体

我们知道,AI产品的研发,从研究到原型再到生产,是一个非常复杂的流程,这造成很多问题,比如精度、速度不及预期。

 

传统的深度学习流程中,训练框架和推理框架往往分开,设计一个算法时,首先经由训练框架,成为一个可训练模型,再转换成推理框架可接受的表示,并在不同设备上进行计算。

 

这里会有一个训练和推理的转换流程,以及中间的模型损失。比如有些算子不被支持,导致无法自动完成转换,甚至转换过程引入大量冗余算子,造成模型性能不理想;以及最后投放芯片时,发现了问题,却又因流程过于复杂无法回溯问题。

 

天元「训推一体」图示

 

面对这种「从研究到生产流程复杂,精度对不齐?」的问题,天元推出「训练推理一体」的方案,即训练推理可以同时进行,它有四个特点:

 

  1. 无需模型转换,直接使用训练模型进行推理;

  2. 速度&精度保持一致,保证训练实现的速度和精度,在推理时保持一致;

  3. 跨设备模型精度对齐,模型训练结束之后,用于不同设备时,保证跨设备模型精度对齐;

  4. 自动模型优化简化流程,天元内置自动模型优化模块,减少手工模型遇错处理,简化整个流程,以实现高效研发体系。

动静合一

深度学习框架大致分为两类,一类是静态框架,它易于部署,快速实现产品,为工业界所喜爱,但是难调适;另一类是动态框架,调适更方便,使用更灵活,为学术界所接受,但资源占用多。这是一个典型的鱼与熊掌的问题。

 

面对「静态图好部署,动态图易调试,二者难兼得?」这一困境,天元希望达成一种「动静合一」的效果。

 

天元「动静合一」代码示例

 

天元通过一个Python装饰器——「@trace」——装饰一段函数,使其即可在动态下正确运行,又可切换到静态形态,你只需把「Enabled」开关设为True或者False,如上图代码所示。

 

由此,我们既可以在动态过程中,方便地做原型的研发调试,又可以同时在生产环节做提速,加速效果可达到5%~20%。

兼容并包

深度学习发展至今,框架众多,接口各异,如果开发者要复现模型,需要学习多种框架知识,方可跟上深度学习最前沿,这无疑是一个巨大的学习成本。面对「框架接口各异,模型复现困难,学习成本高?」的问题,天元发展成了一个「兼容并包」的体系。

 

天元「兼容并包」图示

 

天元的深度学习代码,写法很像Numpy和PyTorch,这种整体简化的Pythonic风格,易于为Python用户所接受;而且,函数的命名风格和参数设计细节,也尊重Python社区原有传统。值得一提的是,天元还有一个实验性功能,支持直接导入PyTorch Module,方便模型复现、原型系统开发和实验。

 

另外,旷视在CV领域的独特积累也融入了天元,集成了众多专为CV优化的算子,使得CV研发更简便。这一切都使得天元系统更加兼容并包。

灵活高效

众所周知,一家AI公司通常面临着非常多的计算设备和场景,如果要提供极致的AI产品和服务,就要在每种设备上打造极致的性能,这个挑战非常大。面对「生产环境计算设备繁多,却需要极致性能?」的问题,天元秉持「灵活高效」的原则,借助高性能算子库,充分释放不同设备的算力。

 

这里是天元与参考框架的训练性能对比,可以看到在主流模型上,天元的速度表现非常不俗;同时在推理性能对比上(CPU场景),天元也保持高性能,提升效果显著。

 

另外,内存的节省是深度学习系统在跨设备部署或者大模型训练时所特别关注的。天元内置了高效的内存优化策略,可以有效减少训练的显存占用,在相同设备上支持更大模型、更多算法。

 

天元「灵活高效」图示

 

此外,天元还有强大的自动内存/速度优化机制,比如上述的亚线性内存优化,在计算速度没有降低的前提下,达到256Batch的训练能力,并且模型越深,效果越好;内部评测上,它实现了20+倍的内存节省,速度几乎不变。

 

为支持更多计算设备,发掘其计算能力,天元引入了JIT自动代码生成技术,显著提升计算速度;内置了自动算法选择机制,在不同设备上自动调优,找到最适宜该设备的算法实现,达到灵活高效的目标。

 

综上所述,天元有着「训练推理一体」、「动静合一」、「兼容并包」、「灵活高效」四个优势。最后,经过打磨,天元还支持顶级研究团队的最前沿学术开发。我们相信,以天元深度学习框架为基础,可以打造新一代AI生产力平台,让开发者更简单、更高效地进行AI生产力的研发。

天元的开发路线

 

天元这次发布是Alpha版,6月份会发布Beta版,天元将支持ARM系列CPU,和更多的加速设备,提供量化和低比特计算支持,推动嵌入式的研发工作;9月份发布正式版,天元会全面覆盖主流计算设备,升级现有动态计算能力,让动态优化和静态优化一样好,带来「训推一体」的更佳体验。

天元研MegStudio视频演示

天元具备MegStudio的在线试用环境,提供免费算力、公开数据集,让你快速体验、学习和开发深度学习项目。同时欢迎大家访问天元官网(megengine.org.cn)和Brain++官网(brainpp.megvii.com),我们会在天元ModelHub中持续更新SOTA算法。

 

开发人员都知道,一个代码自己用,和拿出来与人分享,这之间的差距有多大。天元是一个开诚布公的开源项目,我们希望在正式版发布之前,有更多的人提出建议,给出批评,一起贡献code,一起来共建更好的天元深度学习学习框架,用AI的力量,创造更多的价值。

共创深度学习框架新纪元

旷视最早是做深度学习算法起家,我们希望把算法和应用场景结合起来,为客户提供一整套解决方案,这就需要数字化中台软件和硬件中台的支撑。我们会发现,算法和场景都是一个无限游戏。场景的无限性要求算法的无限性,而要满足算法的无限性,则需要强大的AI基础设施。

 

旷视Brain++正是这样一个AI基础设施,它具备三位一体的深度学习框架天元MegEngine,深度学习云计算平台MegCompute和数据管理平台MegData,合而为一成为新一代AI生产力平台,可使算法生产自动化、规模化和高效化。而这次开源的天元,是提高AI基础设施算法生产效能的核心推动力;更好地帮助开发者进行创造,助力深度学习简单开发,共创深度学习框架新纪元。

天元官网:https://megengine.org.cn

天元GitHub:https://github.com/MegEngine

天元科普视频

传送门 

欢迎大家关注如下 旷视研究院 官方微信号????

发布了55 篇原创文章 · 获赞 3 · 访问量 2万+
展开阅读全文

没有更多推荐了,返回首页

分享到微信朋友圈

×

扫一扫,手机浏览