R TALK | 旷视研究院田忠博：天元助力「深度学习简单开发」 - 极术社区

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

屏幕快照 2020-04-09 下午8.07.02.png

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

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

天元的整体架构

屏幕快照 2020-04-09 下午8.07.10.png

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

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

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

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

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

天元的四个优势

优势1：训练推理一体

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

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

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

无需进行模型转换，直接使用训练后的模型进行推理；
可以保证训练时实现的速度和精度，能够在推理时是保持一致的；
在不同的设备上进行推理和使用时，能够保证跨设备的模型精度实现对齐；
内置自动模型优化过程，大幅减少手工模型遇错处理的机会，把整个流程简化，形成一个非常高效的研发体系。

优势2：动静合一

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

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

屏幕快照 2020-04-09 下午8.07.38.png

这里展示的是天元（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以后，计算速度会有显著提升，获得非常好的效率。另外，天元还内置了自动算法选择机制，在不同的设备上自动进行调优，找到适用于这个设备的最好的算法实现。

天元的开发路线

屏幕快照 2020-04-09 下午8.08.17.png

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

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

视频链接：https://v.qq.com/x/page/t0938jmaela.html

新一代AI生产力平台

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

屏幕快照 2020-04-09 下午8.08.33.png