训推一体，动静合一：深度学习框架「天元」出世

Posted 2021-07-19 机器之心

机器之心报道

机器之心编辑部

「深度学习，简单开发」，刚刚开源的天元（MegEngine）将为开发者们带来全新的深度学习开发体验。

又一个国产深度学习框架开源了。3 月 25 日，旷视科技在北京发布了新一代人工智能生产平台天元（MegEngine），其「训练推理一体」、「动静合一」、「兼容并包」、「灵活高效」的强大性能或许将为 AI 领域带来新的风向。

「天元是旷视 1400 名研发人员一直以来使用的深度学习引擎，已部署在上百个产品、几十种计算平台上，」旷视首席科学家、研究院长孙剑在发布会上表示。「他可以帮助开发者做到协同设计训得好，高效系统训得快，大规模数据或模型训得动。」

旷视首席科学家、研究院长孙剑发布会现场分享。

旷视提出的 AI 生产力平台 Brain++包括深度学习框架（天元 MegEngine）、深度学习云计算平台（MegCompute）、以及数据管理平台（MegData），将算法、算力和数据能力集为一体。本次开源的天元是 Brain++平台的核心。

• 天元 MegEngine 深度学习框架官方网站：https://megengine.org.cn/

时至 2020，自贾扬清等人开放 Caffe 已过去六年，旷视为什么选择在这个时间点开源自身的核心技术？旷视研究院高级技术总监、天元项目负责人田忠博告诉我们：「这一进程其实去年已经开始推动，我们认为目前国内市场上缺少天元这样的项目。」

这是一个需要勇气的行动。据了解，在做出开源的决定以前，旷视技术团队内部经过了多次讨论。这关系到框架技术是否成熟，生态环境是否发展完善，以及未来发展的方向。

田忠博指出，天元是一套训练推理一体化、动静态合一的工业级深度学习框架，具体架构包括五层：

从最上层的接口层到最底层的计算内核，天元做了很多优化，例如内存管理中的亚线性优化和高性能算子库等。

天元架构是旷视在研发过程中遇到痛点，并逐步解决的成果。旷视相信这些痛点也是行业共通的难题，通过这一开源架构，我们可以一站式地解决这些问题。

早在 2014 年 8 月，深度学习框架天元就在旷视内部诞生了，它随后成为了 Brain++的核心组成部分。「自 2012 年起，陆续有研究者开发出了第一代深度学习框架。在旷视，我们在实践过程中发现手工编写每层配置文件的方式非常不灵活，」孙剑介绍道。「因此我们发展了第二代深度学习框架，通过计算图的方式设计模型并求导，大大提高了效率。2015 年谷歌正式发布了 TensorFlow，采用了和旷视殊途同归的思路。当时我们进行了比较，发现自己的方法更具优势，所以天元就继续发展下来了。」

天元与目前流行的 TensorFlow、PyTorch 相似，旨在让更多的人可以在其基础上构建自己的应用，创造属于自己的价值。

此前，百度等公司也推出了自己的深度学习框架。天元会有哪些不同？「在这其中，我们是唯一一个以 AI 为本业的公司，天然需要使用深度学习算法解决问题，」旷视副总裁谢忆楠表示，「我们把自己在 AI 解决各领域问题过程中的工具开源出来，从视角上看或许会获得更好的效果。」

天元跻身全球主流框架

机器之心很多读者都用过深度学习框架，很多读者也都想入门深度学习框架。那么假设我们有 Python 和 NumPy 数值计算方面的基础，旷视新开源的天元又会有哪些方面吸引我们？

孙剑表示，天元有三大特别吸引人的优点，「第一是框架与算法的协同，框架需要为硬件开发优化算子，这样才能结合框架、设备开发最好的网络，ShuffleNet 系列高效神经网络就是协同设计的结果之一。第二即高效与强大性能，一般训练 COCO 数据集需要 30 多个小时，2017 年基于天元的 MegDet 并行物体检测系统，2 小时就能完成，并且精度更高。」

「第三是很好地支持大数据、大规模训练。借天元开源之际，我们同时也发布最新版 Objects365 V2，它是世界上最大的物体检测数据集，是目前学界流行的 COCO 数据集 的 16 倍大小。有这么大规模的数据，模型也需要扩增。如果模型扩大 10 倍的话，总体训练量就会增加 160 倍。而天元对于这种大模型、大数据也是没问题的，」孙剑接着说。

田忠博说：「天元是为工业级、研究院规模的研发机构设计的，它具备的一些核心特点都是用来解决现实研发过程中的痛点。」从具体体验与特性上，我们确实能发现天元解决很多开发痛点，它们从代码风格、编程范式到底层优化展现了不一样的思路。

一键切换训练推理

在机器学习中，训练与推理一直不是一体的，当我们设计算法时，首先需要训练框架的支持，然后再把训练好的模型转移到推理框架上，从而将后续的推理计算放到各种硬件上。这中间会存在训练与推理的转换，这个过程需要手动完成，甚至转换后也不能保证模型性能与进度。

「训练推理一体化能把繁杂的模型生产流程简化到一步、两步，我们也希望它能够极大缩短算法研发和部署的时间，真正支持大规模的生产级诉求」，田忠博说，「为此，天元框架需要做到三点：无需模型转换就能做推理；训练的速度、精度与推理保持一致；在不同设备上做推理需要保证模型精度可以对齐。」

从算法直接到部署硬件，中间不再需要转换模型，只需要几行代码与 API，模型的实践效果就能看得到。 对于大多数算法工程师，可能比较精通深度学习模型，对推理与模型部署相对而言不是那么熟，天元无疑能降低大量工程实践成本。

更重要的是，在部署过程中，天元通过自动算法选择机制和一系列针对推理的自动优化，能帮我们删除冗余代码，找到最适宜该设备的算法实现，达到灵活高效的目标。

动静合一，图并存

从 Theano 到 TensorFlow 1.X，它们都遵循着「静态计算图」这条路，我们会发现在用这类方法写模型时，需要很强的「向量化思维」，For 循环之类的语句就不要多想了。此外，因为写代码只是建计算图，调试的时候异常困难，我们不知道图中的数据流到底是什么样的。

很快，在 PyTorch 等框架中出现了动态计算图，用它写模型就像写个普通 Python 脚本，但这样写出来的模型，却是难以高效部署的。天元同时支持动态图与静态图，在科研训练与调试时采用动态图，到部署的时候又切换到静态图，这样就能充分利用两者优势。

如上图所示，天元可以直接通过 @trace 装饰器确定某个函数到底是动态图还是静态图，而且状态还可以手动调整。这不就意味着在调试时我们可以将函数设置为动态图，直接查看每个变量的信息，而要到了部署等实际应用时，再把静态图打开，提升训练与推断速度？

如果以这样的方式调用动态与静态计算图，那真的称得上整合了两种方式的优势。

此外，我们还发现有一种动、静态的「混合编程」，即在动态图中调用静态图。可以想象，定义网络结构等过程使用动态图比较直观，而后采用静态图的方式执行实际运算，这也是非常吸引人的方式。后文在代码示例中，也会展示这种新颖的「混合编程」。

人生苦短，Pythonic 好入门

如果深度学习建模一直采用静态计算图，估计框架就会劝退很多开发者。但是天元这种主要采用动态图的框架，构建模型就像搭积木，而且还能以一般的 Python 编程方式去搭。相信很多接触机器学习的开发者都会选择 Python，它加上 NumPy 等数值计算库，差不多就能完成大多数任务。

如果有了这些基础，再搭建深度模型就没问题了，天元一样承接 Pythonic API，学习成本非常低。此外，因为旷视本身在计算机视觉方面有非常雄厚的累积，田忠博表示：「天元内置了一些独特的深度学习算子，它们和旷视在 CV 上的独到能力息息相关，能针对 CV 模型的训练与应用实现深度优化。」

对于 API 与代码风格，我们还是比较推荐直接查看本文后面的示例，可以直观感受到，不论把天元当做深度学习的入门框架，还是与 PyTorch、TensorFlow 等一起作为备选框架都没问题。天元的编程风格与 API 配置，和已有的编程习惯、现有框架的风格都是一致的。

当然，这些并不是天元所有特性，还有很多需要进一步探索。例如它还有很强的多平台、多设备适应能力，既能灵活调用硬件能力，又可以高效训练大模型。

天元框架实际体验怎么样？

那么实际用天元写模型是什么样的体验？机器之心体验了旷视发布的 MegStudio，通过在线开发的方式，我们可以快体验到天元的特性与编程方式。

MegStudio 上准备好了一些入门项目，它展示了新框架的使用方式，例如「基本概念」主要介绍什么是计算图、张量和算子等基础知识，每一个项目都是通过中文文字与可运行的代码组成的，就像常用的 JupyterLab 那样，通过运行这些基本项目可以让开发者快速掌握天元的基本使用。

为了直接展示天元的风格，我们可以快速试试 DL 入门的「第一个模型」MNIST 图像识别。准备数据这里就不展示了，天元可以直接调用 API 来处理常见数据，准备 MNIST 数据两行代码就够了。如果要搭建一个简单的卷积神经网络，天元允许使用 PyTorch Module 的方式定义，搭建模型就是搭积木。

如上所示，继承 M.Module 后，Net 类只需要定义初始化参数与重写前馈 forward 过程就可以了。卷积层、批归一化层、激活函数等模块的使用非常符合直接，只需要给定必要参数就没问题了。

这种继承 Module 的方式非常直观，是一种动态计算图编写方式。如果要实现前向传播，传给它输入数据就行了，要是需要调试，在 forward 函数下设置断点，或者打印变量的信息都没问题。

然而，一般写模型真正复杂的地方在于定义模型架构之外的过程，例如后面需要完成的训练与测试过程。天元的优势在于，它允许动态图结合静态图，即使定义模型使用的是动态图，但在定义训练与测试函数仍然可以用静态计算图。

如上在最为核心的前向、反向传播与模型损失计算过程中，我们可以采用天元中的 @trace 装饰器，直接调用静态计算图模式。因为绝大多数计算都发生在这个过程，因此，由动态计算图转为静态计算图将更有优势，天元能采用更高效的方式训练模型。

最后就是迭代训练过程了，这里天元也采用了 DataLoader，这种方式非常便捷。后续的损失函数计算与最优化器更新，与现有的 DL 框架比较类似，我们差不多不需要学习额外的知识，了解一下 API 就能用。

当然，后续还有保存、加载模型，也是两条语句的功夫。总体而言，很明显天元类似 PyTorch 保持了最好理解与使用的接口，并在关键的计算过程提供装饰器以调用静态计算图，这可以称得上是结合了动态、静态图各自的优势。

因为整体上天元的 API 接口非常精简，且保持着主流框架所养成的编程习惯，入门门槛还是非常低的，之前了解过深度学习框架的同学，不需要文档都能快速读懂整个过程。此外天元作为新手入门框架也没是没问题的，学会了它，差不多等同于学会了深度学习常用框架。

此外，除了 MegStudio，旷视还发布了模型中心 ModelHub，它汇聚了全球顶尖算法的一些预训练的模型，和旷视研究院一些最新的技术、研发成果。

旷视联合创始人兼 CTO 唐文斌 发布会现场分享

在旷视的计划中，Brain++将作为「AI 基础设施」逐步开源、开放，将生产力共享给广大开发者，让每个拥有小场景需求的人都可以做出能够解决其关键问题的最优算法。此外，旷视还希望在未来面向高校课程、初创企业提供更多支持，以方便小规模深度学习部署。

值得一提的是，旷视本次发布的天元版本是 Alpha 版。旷视表示，更为「正式」的版本会在今年晚些时候上线，随着大量组件的更新，旷视会逐步将计划中开源的代码和能力进一步释放出来。「3 月份开源的版本包含前期整理的代码和关键组件，今年 6 月我们还会发布对 Arm 架构、深度学习加速芯片、低比特和量化环境的支持。在今年 9 月份，天元还会有动态图的核心升级，整个系统会走向第一个开源的稳定版。届时体验会达到更为理想的状态。」田忠博表示。

自 2012 年起，旷视就推出了人脸识别开放平台 Face++，近几年来，AI 已从实验室中快速走向实用化。旷视的研发人员们认为，「天元」的开源宣告着人工智能技术再次回到了算法探索的层面。

未来，旷视还会进一步增强这一框架的分布式机器学习能力，对底层编译与代码生成能力进行优化，并欢迎机器学习社区力量的加入。

本文为机器之心报道，转载请联系本公众号获得授权。

✄------------------------------------------------

加入机器之心（全职记者 / 实习生）：[e-mail]

投稿或寻求报道：content @jiqizhixin.com

广告 & 商务合作：[e-mail]