35 万行代码，旷视重磅开源天元深度学习框架，四大特性实现简单开发

110 阅读 0 评论 73 点赞

我是靠谱客的博主生动鸵鸟，最近开发中收集的这篇文章主要介绍35 万行代码，旷视重磅开源天元深度学习框架，四大特性实现简单开发，觉得挺不错的，现在分享给大家，希望可以做个参考。

概述

【导读】2020 年 3 月 25 日，人工智能企业旷视科技举办线上发布会，旷视联合创始人兼 CTO 唐文斌宣布正式开源其 AI 生产力平台 Brain++ 的核心组件——天元(MegEngine)。本次发布为 Alpha 版本，基于 Apache License 2.0，向外界共开源约 35 万行代码，包括 C++、CUDA 和 Python代码，在GitHub 上进行发布。

发布会上，旷视研究院高级技术总监田忠博详细介绍了这款刚刚正式对外开源的深度学习框架。

全球AI开发框架又添一员，旷视开源“天元”

旷视天元开源之时，正值深度学习框架百花齐放的时代。

自 2007 年 Theano 诞生以来，经过十余年发展，深度学习技术与应用突飞猛进，深度学习框架也处在不断迭代与进化的过程；另一方面，开源的概念在全球范围内越来越深入人心，这使得人工智能开发依赖的环境安装、部署、测试，以及不断迭代改进准确性和性能调优的工作变得更加简单，在人工智能领域，开源深度学习框架已经成为开发者离不开的平台和工具。

学界和业界的共同努力下，诞生了早期从学术界走出的 Caffe、 Torch 和 Theano，现如今引领着产业界的 TensorFlow，Amazon 押注的 MXNet，Facebook 倾力打造的 PyTorch，Microsoft 内部开源的 CNTK，以及相对小众的深度学习引擎 DSSTNE 等深度学习框架。

简单梳理一下这些主流深度学习框架发展的历程，我们会发现它们各有各的特性：

图片已获授权

TensorFlow 由 Google 于 2015 年 11 月正式开源，很快就成为深度学习领域占据绝对统治地位的深度学习框架，很多企业的产品都基于这一框架开发，如小米、京东、Airbnb 等。TensorFlow 全面的开发语言和模型训练服务器、移动设备支持，使得其成为产业界采用最多的深度学习框架。

MXNet 项目诞生于 2015 年 9 月，当时在卡耐基梅隆大学 CMU 读博的李沐创造了这个轻量级、可移植、灵活的分布式的开源深度学习框架，后成为 Amazon 官方主推的深度学习框架，支持 CNN、RNN、LTSM，为图像、手写文字和语音的识别和预测以及自然语言处理提供了出色的工具。

Keras 的创造者是谷歌 AI 研究员 Francois Chollet，自 2015 年 11 月开源以来，已发展成为第二大流行深度学习框架。这个由 Python 编写的开源人工神经网络库可以作为 Tensorflow 、CNTK 和 Theano 的高阶应用程序接口，进行深度学习模型的设计、调试、评估、应用和可视化，目标是只需几行代码就能让你构建一个神经网络。

2016 年，微软开发的认知工具包 CNTK 问世，支持 RNN 和 CNN 类型的神经模型，成为处理图像，手写和语音识别问题的最佳候选者。虽然 CNTK 的分布式计算性能较高，但缺乏对 ARM 架构的支持限制了其在移动设备上的功能。

2017 年，Facebook 开源了用于神经网络训练的 Python 包 PyTorch，它改编自基于 Lua 的深度学习库 Torch，类似于 Numpy，非常 Python 化，很容易就能和 Python 生态系统的其他部分集成。由于对动态图的支持， PyTorch 的灵活性相比 TensorFlow 大大提升，特别是用于快速验证和算法复现，因此备受学术界的青睐。

有了这些功能强大的开发框架，AI 开发者基本上也都会用之进行科研或业务落地。但是在人工智能领域，大家使用比较多的还是 Google、Facebook、微软、亚马逊的开源框架，国内虽然有很多互联网巨头都在开始这方面的工作，但目前还没有形成风潮。

2016 年，互联网巨头百度开源了飞桨（PaddlePaddle），可能是国内目前最有影响力的 AI 框架；2019 年，通讯行业巨头华为宣布即将在 2020 年一季度开源 MindSpore，但目前仍无进一步消息；3 月 25日，旷视研发的深度学习框架天元（MegEngine）正式开源。

与主流深度学习框架相比，旷视的 MegEngine 有哪些特点呢？

开源 35 万行代码，天元技术架构理念新颖

唐文斌介绍，本次旷视天元共开源约 35 万行代码，包括 C++、CUDA 和 Python 的代码。

旷视联合创始人兼 CTO 唐文斌

天元是一套伴随旷视自身 AI 产业实战经验的框架，是旷视 Brain++ 的核心组件之一。为了这次开源，旷视为天元做了一次全面的升级。

从 2014 年开始研发，2015 年全员使用，到今年 3 月开源，旷视目前所有的算法都是基于天元 MegEngine 这个框架训练和推理的。它不仅能够在 AI竞赛擂台上为旷视打怪升级加 Buff，更撑起了旷视工程化、产品化的半边天。

发布会上，天元项目的负责人，也是旷视研究院高级技术总监田忠博指出，天元是一套训练推理一体化、动静态合一的工业级深度学习框架。

从上到下，天元可以分为五个层次，最上面是计算接口层，向外连接了 Python 和 C++ 接口，开发者可以通过 Python 和 C++ 两种语言对整个框架进行使用和编程，以及系统的设计和研发、训练和推理。

接着是图表示层，包含了动态图和静态图的表示功能。

再往下是一个完整的一体化核心计算引擎，具有自动求导机制，图优化和图编译功能，有了这个层次就可以支撑起动态、静态和接口完整的功能。

在这个层次之下的运行时管理层由两个主要部分组成，一部分是计算调度，可以将计算设备抽象为执行流，由调度器对这些执行流进行合理的调度；另一部分是一整套内存管理机制，包括静态内存和动态内存管理。此外，这个模块里还内置了许多关于内存的高级优化，其中值得一提的是，在其中实现了静态的亚线性内存的优化器，使得内存管理效率得到大幅提升。

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

四大特性剖析天元如何实现简单开发

旷视研究院高级技术总监田忠博

在过去几年，旷视在研发过程中遇到了很多行业共通的痛点，而天元的核心特色就是紧紧围绕着这些痛点的。

具有来说，天元四大核心特性：训练推理一体化、动静合一、兼容并包和灵活高效。

训练推理一体化

比如其中的一个痛点，是深度学习从研究到生产的流程非常复杂，各个阶段模型精度往往很难对齐。

田忠博指出，在传统深度学习研发流程中，训练框架和推理框架往往会分别设计和实现，训练框架和推理框架是两个阶段，当进行算法设计时，这个算法要首先经过训练框架的支持，变成一个可训练的模型，还要再把它转换到一个推理框架上可以接受新的表示，再由推理框架在不同的计算设备上进行计算。

在这里会有一个训练和推理的转换过程，这一过程中会产生很多问题，比如因为训练框架和推理框架是分别设计的，所以其中有些算力可能不被支持，导致无法自动完成转换，需要手工进行优化，转换过程中也可能引入了大量冗余的算子，致使最后的模型性能和精度并不理想。当最后把推理框架投放在芯片上进行计算时问题暴露，但因为整个流程复杂，我们无法精准地找到问题所在。

因此，天元框架的设计理念，就是希望训练和推理一体，即让它能够同时进行训练，也能够进行推理。

针对这个痛点，天元的训练推理一体化可以很好地解决。

（1）它无需进行模型的转换，可以直接使用训练后得到的模型进行推理；

（2）可以通过这一机制，保证训练的速度和精度与推理保持一致；

（3）模型训练结束后，需要在不同的设备上进行推理、使用，该框架也能够保证跨设备的模型精度实现对齐（最小化精度差别）；

（4）通过简化流程，天元框架能够内置一个自动模型优化过程，减少手工模型遇错处理，可以直接自动使用内置流程，简化流程，形成高效的研发体系。

这样一来，AI 真正落地要考虑的多端部署和在线服务的问题就得以解决，大大减少了训练成本的问题。

动静合一

痛点二，静态图好部署，动态图易调试，但二者难以兼得。田忠博介绍道，深度学习框架大致分为两类，一类是以 TensorFlow 1.0 为代表的静态深度学习框架，它非常容易部署，能够很快地产出产品，是现在工业界非常喜欢的部署方式，它的性能高，占用资源少，但是难以调试。在学界，大家更喜欢以 PyTorch 为代表的动态计算框架，因为它在研究阶段调试更方便，使用更灵活，但是动态图也有缺陷，比如内存占用严重，很难做优化等。

面对这个鱼与熊掌不可兼得的问题，旷视尝试把两种框架的优点集成在一起，在设计天元时希望能够达到动静合一的效果。

上图展示的是天元框架代码中从动态到静态切换的情况。可以看到，通过使用一个 @trace 的 Python 装饰器来装饰其中一段函数，实现了这段函数具备既可在动态下正确运行，也可以转换到静态形态运行的状态。只需把「Enabled」开关设为True或者False，用户就可以自由选择动态或静态计算。

这样，开发者就可以在动态的过程中，非常方便地进行原型的研发和调试，同时当希望在生产环节使用，或希望借助更好的静态优化器、静态编译机制进行提速时，可借助静态图进行提速。

田忠博表示，在测试中，静态提速往往可以达到 5% 到 20% 的加速效果，节省时间，提高效率。

兼容并包

第三个痛点，是市面上有很多框架，但每种框架使用的接口都不一样，这导致大家在进行学术交流时，首先要了解它是用什么框架实现的，在使用中还需要在常用的环境和框架中再重新进行模型实现，这对于一般的开发者来讲是一件高成本的事。

因此，为了简化这个问题，天元在设计时还希望它是一个兼容并包的体系。

上图为使用天元框架进行深度学习的代码，它的风格与Numpy 和 PyTorch 的写法非常相似，Pythonic 风格的简化 API 让 Python 使用者可以自然地接受，所以在函数的命名风格和参数的设计细节中尊重原有 Python 社区的传统。

值得一提的是，天元还提供一个实验性的功能，让开发者可以便利地将以往写过的模块，如将PyTorch Module 直接导入到框架中，和其他天元组件一起使用，以更好地进行模型复现。

另外，田忠博提到，旷视在计算机视觉领域有一些独特的积累，因此也把其在这方面的成果融入到天元系统中，集成了很多为计算机视觉特别优化的算子，让计算机视觉研发更加简便。

灵活高效

痛点四，对于一家进 AI 生产公司来说，可能会面临很多设备和场景，需要在每一种设备上实现极致的性能。

在框架设计时，天元秉持要灵活高效的原则，在许多的设备、算法上，都能得到领先的性能。接下来，田忠博放出了训练性能对比图，与若干擅长推理的框架进行横向对比。

结果显示，在 CPU 推理场景下，天元在训练性能上有显著的提升和优势，即可以同时在训练和推理过程中保持高性能。另外，如果要把算法更好地部署在各种设备中，或者在训练时能够利用现有的设备训练更大的模型，支持更多的算法种类，显存或设备的片上内存使用是一个非常关键的因素。所以，节省内存也是天元所关注的。

天元内置了一个高效的内存优化策略，它能够显著减少训练时的显存占用，实现在同样的设备上可以训练更大的模型，支持更多算法。

此外，天元还有很多内存和速度的优化机制，比如亚线性内存优化。可以发现，在使用天元动态图能力时，可以支持 32 Batch 左右的计算；如果转换到静态图下，就可以支持 64 Batch 的计算。那么，如果希望在这种情况下，训练更大的 Batch 和模型，则完全可以在这里采用亚线性自动内存优化技术，在几乎不降低计算速度前提下，达到 256 Batch 的训练能力，而且模型越大、越深，它的效果越好。

田忠博表示，在内部评测中，天元可以实现某些大模型训练时内存节省 20 倍以上，而速度几乎不变。

这些特性，使得天元能够实现产品从实验室原型到工业上能够部署的小时级转化能力，以及大规模的、弹性的训练，并支撑顶级研究团队进行最前沿的学术开发。

这样，天元可以做到“简单开发”，让开发者真正体验到“训得好”、“训得动”、“训得快”。

揭秘天元“前世今生”，研发路线图首次曝光

从 Theano 为源头，到不断迭代到今天发布的MegEngine Alpha 版本，天元的诞生来之不易，背后是旷视研究院团队从 0 到 1 的打磨过程。

旷视成立初衷是希望把计算机视觉应用于传统产业，用技术改变世界。当 2013 年中深度学习刚刚兴起之时，清华宿舍中一名实习生埋首两周，研发出一套人脸识别检测算法，算法性能技惊四座，于是旷视正式走上用神经网络解决一切问题的道路。

起初，旷视用 Theano 框架写模型代码，训练神经网络，但随着网络越训越大，越来越复杂，低效耗时的框架令人崩溃，公司中的一些大牛开始琢磨其他的办法。

2013 年底，旷视当时的研发负责人曹志敏提出打造一套能够打通数据、训练和业务的自动化算法研发系统 Cycle++，不需要投入过多人力和时间就可以实现算法从研发到应用的自循环体系（旷视 Brain++的早期设想）。于是，2014 年初，旷视自研的初版深度学习框架诞生了。

经过磨合，旷视在 2015 年年中完成了自研框架与公司内部所有业务的接轨，公司业务线上的模型全部换成了自研框架训练出来的版本。

2015 年 11 月 9 日，Google 正式发布并开源 TensorFlow，旷视发现原来他们是殊途同归，都是基于计算图的方式来做框架，但这也给旷视的自研框架造成很大冲击，公司内部在是否要继续坚持自研框架上发生分歧。经过激烈的讨论和详细的评测后，旷视发现当时的TensorFlow 的性能并不理想，竟比自研框架要慢若干倍。最终旷视选择了坚持自研的道路。

此后，经过不断迭代，同时在工业实践的锻炼中，除了最底层的框架，旷视也在同时进行数据和算力基础设施的变革。2013 年，旷视研究院成立了自己的数据团队，随着业务数据的暴增，数据管理不断出现问题，旷视又开始建立自己的数据管理系统 MegData。

2015 年底，天元MegEngine 已经进入了稳步发展期，但公司“小作坊”模式开始扛不住业务需求，计算资源成为瓶颈问题，于是旷视建设了“正经的机房”，研发出深度学习云计算平台 MegCompute，并仅用一个季度的时间完成了业务从单机到集群的彻底迁移。

旷视从研发到业务全面向自有深度学习框架和自有计算集群的迁移，标志着旷视数据、算法和算力三个核心组件正式完成“大一统”，自此旷视 AI 生产力平台 Brain++ 雏形初现。

2016 年，旷视开始组建大规模的团队持续优化 Brain++ 的整个套件开发流程，2019 年开始筹备将 Brain++ 最核心的深度学习框架开源，并为 MegEngine 起了一个中文名字——天元。这期间框架研发团队可以说是经历了一场浴火重生，把原来封装好的代码分解再重组，让开发者上手更快。

经过一年的筹备，天元今天终于如期开源，赋能开发者。未来，天元还有更多计划，发布会现场旷视就首次曝光了天元的开发路线图。

田忠博表示，本次旷视开源的天元是 Alpha 版本，未来的开发计划是在今年 6 月份发布 Beta 版本，届时天元将提供ARM 系列 CPU 支持，更多的加速设备支持，以及量化和低比特计算支持；到 9 月份发布正式 1.0 版本时，天元支持的主流计算设备将更全面，动态能力升级，并优化训练推理全流程使用体验。

他说，在 Beta 版本和正式版本之间，希望更多人能够参与并贡献 code，“也许下一代天元并不是由旷视的研发团队做出来的，而是与你一起共创出来的 Beta 和正式版本，所以我们也希望跟大家一起来共建更好的深度学习框架。”