R TALK | 旷视研究院田忠博:天元助力「深度学习简单开发」


「R TALK」是一个深度学习专栏,由北京智源-旷视联合实验室推出,旨在通过一场场精彩纷呈的深度学习演讲,展示旷视研究院的学术分享及阶段性技术成果,抛砖引玉,推陈出新,推动中国乃至全球领域深度学习技术的发展。这是「R Talk」第 12 期分享。

大家好,我是田忠博,非常荣幸有机会向大家介绍天元(MegEngine)深度学习框架。

 

天元(MegEngine)是旷视完全自主研发的深度学习框架,它的研发始于2014年,到现在已有6个年头;旷视全面使用天元构建了几乎所有的AI产品。针对这次开源(注:3月25日,旷视召开Brain++ • 天元深度学习框架开源发布会,宣布「天元」开源),我们特别做了大型升级,让天元以全新面貌面对大家。

天元的整体架构

 

这张图展示的是天元(MegEngine)深度学习框架的整体架构,可以看到,从上到下可以分为五个层次。

 

最上面是计算接口层,它向外暴露了Python和C++的接口,可以通过Python和C++两种语言对整个框架进行使用和编程,进行原型系统的设计和研发,进行训练和推理。它的下面是图表示层,包含了动态图和静态图的表示功能。再下面是优化与编译层,它是一个完整的、一体化的核心计算图,有自动求导机制,以及图优化和图编译,支撑起动态图、静态图和接口的完整功能。

 

再往下是运行时管理层,它有两个主要的部分,计算调度和内存管理;前者会把整个计算设备抽象为执行流,由调度器对这些执行流进行合理的调度,后者是一整套的内存管理机制,包括静态/动态内存管理。此外,这个模块里还内置了许多关于内存的高级优化,其中值得一提的是静态的亚线性内存优化器,可以极大提升整个内存管理效率。

 

最底层是支撑整个系统的计算内核层,它有一个高性能的算子库,支持常见的一系列计算设备,比如X86、CUDA、ARM以及专业计算芯片等;同时它还包含一个高性能的异构通信库,可使整个计算框架在分布式多节点上进行大规模的使用,支撑更大规模的训练。

 

可以看到,天元是一个复杂的系统,那么它的目的是什么?是为了解决旷视多年以来在产品中遇到的许多困境和痛点,由此也诞生了天元的四个优势。

天元的四个优势

优势1:训练推理一体

 

第一个痛点是说,整个AI产品研发有一个非常非常复杂的流程,其中速度和精度可能会出现预想不到的问题。传统的深度学习研发流程中,训练框架和推理框架往往分别设计和实现,在这种情况下,当设计一个算法时,它要首先经过训练框架的支持,变成一个可训练的模型,然后把它转换到一个推理框架上可以接受的新表示,再由推理框架把它在不同的计算设备上进行计算。

 

如果大家经常进行深度学习训练和推理的转换,会发现其中有很多问题,比如由于分别设计,一些算力可能不被支持,导致转换无法自动完成,需要手工进行优化;甚至转换时引入了大量冗余算子,使得最后的模型性能和精度并不理想,最后把推理框架投放在芯片上,由于整个流程过于复杂而无法很好地找到问题。

 

有鉴于此,天元(MegEngine)框架从最初的设计理念上,就希望是训练和推理一体的,即能够进行训练,同时也能进行推理。在训练和推理的过程中,它有这样几个特点:

 

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

  2. 可以保证训练时实现的速度和精度,能够在推理时是保持一致的;

  3. 在不同的设备上进行推理和使用时,能够保证跨设备的模型精度实现对齐;

  4. 内置自动模型优化过程,大幅减少手工模型遇错处理的机会,把整个流程简化,形成一个非常高效的研发体系。 

优势2:动静合一

 

第二个痛点是说,静态图好部署,动态图易调试,二者难以兼得。深度学习框架大致分为两类,一类是以TensorFlow为代表的静态深度学习框架,它非常容易部署,性能高,占用资源少,现在工业界非常喜欢,但是难调试;另一类是以PyTorch 为代表的动态深度学习框架,它在研究阶段调试会更加方便,使用起来更加灵活。

 

面对这个鱼和熊掌的问题,我们问自己这样一个问题,这两种计算方式真的很难兼得吗?能否尝试把两种框架的优点集成在一起,更好地进行产品研发和原型开发?由这个问题出发,我们在设计天元(MegEngine)框架时,希望达到一种动静合一的效果。

这里展示的是天元(MegEngine)框架代码中由动到静的切换情况。可以看到,通过一个叫@trace的Python装饰器来装饰其中一段函数,使得这段函数具备既可在动态下正确运行,也可转换到静态形态运行的一种状态。我们只需打开开关,(Enabled 设为 True 或者 False),便可在两种模式上正确计算。

 

由此,我们可以在动态过程中,非常方便地研发和调试原型,同时在生产环节进行使用时,或者借助更好的静态优化器、静态编译机制进行提速时,可以借助静态图把它进行提速。我们测试时,静态提速可以实现5%-20%的加速效果。 

优势3:兼容并包

第三个痛点是市面上有非常多的框架,每种框架使用的接口不太一样。当进行学术交流时,首先要了解这是什么框架实现的,这样大家需要学习多种框架,才能更多地了解到深度学习行业的最新、最好算法。对于一般的开发者来讲,这是一个重大的学习成本。那么,为了简化这个问题,我们在设计天元(MegEngine)深度学习框架时,希望它是一个兼容并包的体系。

 

可以看到,这是使用天元(MegEngine)框架进行深度学习的代码,它的风格非常像Numpy和PyTorch,这种Pythonic风格的接口,让Python的使用者非常自然地接受;另外,它没有过多新的概念,延续大家惯常的理解。所以,天元在函数的命名风格和参数设计的细节中,都会尊重原有Python社区的传统。

 

值得一提的是,天元还提供一个实验性的功能,可以方便地把已经写过的PyTorch Module直接导入框架,方便模型的复现,更简单地进行原型系统的开发和实验。另外,旷视在计算机视觉领域的独特积累也融入了天元(MegEngine)的系统中,从而集成了很多为计算机视觉特别优化的算子,使计算机视觉研发更加简便。

优势4:灵活高效

痛点四是指一家AI生产公司可能面临着众多的设备,和非常多需要解决的场景。如果要在这每一种设备上都能打造极致的性能,这无疑是一个非常大的挑战。在设计天元时,我们秉持一个原则就是灵活高效,希望在许多的设备、算法上,得到领先的性能,希望借助高性能算子库,充分利用不同设备上的算力。

 

这里展示的是我们在训练时的性能对比,通过横向对比几个参考框架,会发现天元(MegEngine)在主流模型上都有非常不俗的速度表现。同时,因为天元(MegEngine)这个系统既用于训练又用于推理,对比时我们会选取一些推理见长的框架。

 

我们发现,在CPU推理场景下,也能有显著的提升和优势。这说明天元可以同时在训练和推理过程中保持高性能。另外,如果把算法更好地部署在各种设备中,或者利用现有设备训练更大的模型,支持更多的算法种类,我们发现,显存或者设备的片上内存是一个非常关键的因素。所以内存的节省也是天元(MegEngine)这样一个深度学习系统所关注的。

 

天元内置了一个高效的内存优化策略,能够非常好地减少训练时的显存占用,在同样的设备上可以训练更大的模型,支持更多的算法种类。此外,天元还有很多强大的内存和速度的优化机制,比如前面提到的亚线性内存优化。

 

当使用天元动态图能力时,它大概可在32Batch左右进行计算;如果转换到静态图,可以支持到64Batch的水平。那么,如果希望训练更大的Batch,训练更大的模型,则完全可以采用天元的亚线性自动内存优化技术,在计算速度几乎没有降低的前提下,达到256Batch的训练能力;而且这个能力随着模型的变大、变深,效果越好。在天元的内部评测上,可以实现某些大模型20倍以上的内存节省,速度几乎不变。

 

为了支持更多的计算设备,同时充分发掘计算能力,天元还引入了JIT自动代码生成技术加速计算。当打开JIT以后,计算速度会有显著提升,获得非常好的效率。另外,天元还内置了自动算法选择机制,在不同的设备上自动进行调优,找到适用于这个设备的最好的算法实现。

天元的开发路线

 

这里展示了天元(MegEngine)的整体研发路线。这次开源的是天元(MegEngine)深度学习框架Alpha版本。我们计划6月份提供第一个里程碑版本Beta版本,在这个版本中,希望全面提供关于ARM系列CPU的支持能力,同时支持更多的加速设备,提供关于量化和低比特的计算支持,使得大家可以基于天元(MegEngine)框架,进行先进的嵌入式研发。

 

9月份,我们希望天元(MegEngine)的能力全面覆盖主流的计算设备;把现有的动态计算能力进行全面升级,让动态图的优化能力和静态图一样好,实现优化训练和推理的全流程体验。下面,我为大家展示MegStudio的在线试用环境,下面是一段展示。

 

新一代AI生产力平台

通过以上的介绍,我们发现天元(MegEngine)框架的一些特点:训练推理一体化、动静合一、兼容并包、灵活高效。它们能够实现产品从实验室原型到工业部署的小时级转化能力。利用这些特点,天元能够实现大规模的、弹性的训练,支撑我们用更大的数据、更大的模型、更好的算法,提供更好的产品。另外,借助这些特性,以及对它们的不断打磨,天元(MegEngine)框架可以支撑顶级的研究团队进行最前沿的学术开发。

 

同时,欢迎大家访问Brain++官网(Brainpp.megvii.com),以及天元(MegEngine)官网(MegEngine.org.cn),我们会在天元网站中、在ModelHub中持续增加新的先进算法,希望帮助大家更好的进行自己的试用和创造。

 

所以,我们相信,以天元(MegEngine)这样一个深度学习框架为基础,可以打造新一代的AI生产力平台,所有人都可以基于这样的平台,更好、更简单地进行AI生产力的研发。

 

谢谢大家。

  

---------R Talker 介绍-------

 

田忠博,旷视研究院研发总监,目前专注于高性能计算系统构建。曾先后就职于网易,百度,豆瓣。

视频回顾

入群交流

添加旷视研究院小助手

 

输入关键词“开源”进入交流群

 

PPT下载

 

在旷视研究院微信公众号后台回复关键词“天元MegEngine”以获取PPT。

往期回顾

R Talk | 旷视首席科学家孙剑:云、端、芯上的视觉计算

R Talk | 旷视研究院目标检测概述:Beyond RetinaNet and Mask R-CNN

R Talk | 旷视研究院姚聪博士:深度学习时代的文字检测与识别技术

R Talk | 旷视南京研究院魏秀参:细粒度图像分析综述

R Talk | 旷视研究院张祥雨:高效轻量级深度模型的研究与实践

R Talk | 旷视研究院张弛:行人重识别及其应用

R Talk | 旷视研究院危夷晨:不确定性学习在视觉识别中的应用

R Talk | 旷视成都研究院负责人刘帅成:图像对齐及其应用

R Talk | 旷视研究院SLAM组负责人刘骁:三维视觉与机器人

R Talk | 旷视研究院周而进:距离对齐技术在人脸识别中的应用

R Talk | 旷视研究院范浩强&周舒畅:AI计算机摄影的原理、应用与硬件设计

传送门 

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

已标记关键词 清除标记
相关推荐