8次迭代5大升级，3大核心优势，旷视深度学习框架天元1.0详细分析讲解

计算机视觉研究院专栏

作者：Edison_G

AI 技术需要开源、开放。就在刚刚结束的中关村论坛旷视平行论坛中，旷视正式发布了天元 1.0 预览版。

9 月 18 日，中关村论坛旷视平行论坛中，旷视研究院院长、首席科学家孙剑正式发布了开源深度学习框架天元 1.0 预览版，并对天元的技术特性与发展方向进行了深入解读。

旷视 2014 年开始研发深度学习框架天元。目前，旷视 1400 多名研发人员全员使用天元 MegEngine，旷视所有算法均通过天元 MegEngine 进行训练和推理。

MegEngine 简介

MegEngine 是旷视完全自主研发的深度学习框架，中文名为“天元”，是旷视 AI 战略的重要组成部分，负责 AI 三要素（算法，算力，数据）中的“算法”。MegEngine 的研发始于 2014 年，旷视内部全员使用。如今，旷视的所有算法均基于 MegEngine 进行训练和推理。

MegEngine 是工业级的深度学习框架，架构先进，性能优异，移植性强。MegEngine 强调产品化能力，在此基础上保证研发过程的快捷便利。

MegEngine 具有几个特点。一是“训练推理一体”。MegEngine 支持多种硬件平台（ CPU，GPU，ARM ）。不同硬件上的推理框架和 MegEngine 的训练框架无缝衔接。部署时无需做额外的模型转换，速度/精度和训练保持一致，有效解决了 AI 落地中“部署环境和训练环境不同，部署难”的问题。

二是“动静合一”。动态图易调试，静态图好部署。鱼和熊掌如何兼得，是现代深度学习框架的核心诉求。MegEngine 在静态图的基础上，逐渐加入支持完整动态图的功能。在动态模式下加速研发过程，无需改变模型代码一键切换至静态模式下的部署，为科研和算法工程师同时提供便利。

三是“兼容并包”。MegEngine 的顶层 API 基于 Python，采取了类似于 PyTorch 的风格。简单直接，易于上手，便于现有项目进行移植或整合。为更好地帮助学习实践，MegEngine 同时提供了“开箱即用”的在线深度学习工具 MegStudio ，和汇聚了顶尖算法和模型的预训练模型集合 Model Hub 。

四是“灵活高效”。MegEngine 底层的高性能算子库对于不同的硬件架构进行了深度适配和优化，并提供高效的亚线性内存优化策略，对于生产环境繁多的计算设备提供了极致的性能保证。高效易用的分布式训练实现能有效支持富有弹性的大规模训练。

MegEngine 的上述特点使其成为了最适合工业级研发的框架之一。更多特性还在持续开发中，也欢迎更多的开发者加入。

动静合一

瞄准痛点：静态图好部署，动态图易调试，但两者难以兼得

• 同时适配科研实验和生产部署环境

• 内置动静转换

• 动静态混合编程

import megengine.functional as F
from megengine.jit import trace


# import trace之后设置 enabled 属性切换动静态图
trace.enabled = True


# 使用 trace 类装饰网络 forward 的函数
@trace
def train_func(data, label, *, opt, net):
    pred = net(data)
    loss = F.cross_entropy_with_softmax(pred, label)
    opt.backward(loss)
    return pred, loss


# 调用函数训练网络，动静态图一套代码
train_func(data, label, opt=optimizer, net=le_net)

兼容并包

瞄准痛点：框架学习接口各异，模型复现困难，学习成本高

• Pythonic风格API，简单直接，易于上手

• 支持导入PyTorch Module

• 特别为计算机视觉（Computer Vision）任务优化

import megengine as mge
import megengine.functional as F
import megengine.module as M
import numpy as np


# 经典的基于 Module 的网络搭建接口
class LeNet(M.Module):
    def __init__(self):
        super(LeNet, self).__init__()
        self.conv1 = M.Conv2d(1, 6, 5)
        self.relu1 = M.ReLU()
        self.pool1 = M.MaxPool2d(2, 2)
        # 省略部分代码...
        self.classifer = M.Linear(84, 10)


    # 符合 Pythonic 风格的计算流程代码
    def forward(self, x):
        x = self.pool1(self.relu1(self.conv1(x)))
        # 省略部分代码...
        x = self.classifer(x)
        return x

灵活高效

瞄准痛点：生产环境计算设备繁多，缺乏优秀性能

• 高性能算子，充分利用算力

• 高效内存优化策略，支持自动 Sublinear 内存优化

• JIT代码生成机制，加速计算

• 内置算法选择，智能适配设备

训练推理一体

瞄准痛点：从研究到生产，流程复杂，精度难以对齐

• 从训练到推理，无需模型转化，精度损失最小化

• 跨设备模型精度对齐

• 自动模型优化简化流程

from megengine.jit import trace


# 使用 trace 类装饰网络 forward 的函数
@trace
def val_func(x, *, net):
    return net(x)


# 调用trace接口无需运行直接编译网络
val_func.trace(inp, net=net)


# 将编译后的网络进行导出，直接生成可用于部署的序列化文件
val_func.dump('./mnist.mge', arg_names=["data"])

目前官网也支持“视觉类”的模型有八类了：

接下来我们可以详细说说天元1.0的改进！

旷视天元 1.0 预览版，历经 8 次迭代 5 大升级

2020 年 3 月，旷视正式将天元开源，提供给全球开发者使用。开源以来，天元从 3 月的 Alpha 版本升级到 6 月的 Beta 版本、到 9 月的 1.0 版本，期间实现了 8 次迭代。

此次天元 1.0 预览版，相比之前的版本，实现了 5 项重要的技术升级。  

首先，天元 1.0 预览版提供了全新的 Imperative Runtime。天元通过重写动态执行引擎，打破过去几个版本中动态图的限制，解决了一系列资源释放的问题，并大幅提升了动态自由度，让使用 GPU 计算像 NumPy 一样方便自如。

其次，天元 1.0 预览版新增自动代码裁剪功能。在实际的 AI 应用开发中，用户经常面临模型大小的问题。自动代码裁剪功能让用户可以全自动的针对自己网络使用的算子进行代码裁剪，不用手工配置就能最小化推理时的代码体积，极大提升端侧推理的竞争力。

此外，天元 1.0 预览版进行了 10 余项推理侧性能优化，进一步提升了端侧推理性能。同时，天元支持了更多的国产硬件。天元对于主流的一些国产硬件进行了接入，方便在国产 NPU 芯片上进行推理工作。

最后，天元实验性的开发了一套基于 MLIR 的 JIT 引擎，尝试利用 MLIR 这一项非常有前景的方案进行计算图的进一步融合、优化来整体提升深度学习训练和推理的速度。

经过了半年的开源工作，现在天元拥有了完整的功能体系。天元希望能依靠训推一体这一特性，让产业应用可以更快、更高效的落地。让深度学习也可以简单开发。

天元三大核心优势，实现简单开发

从开源到现在经过半年的技术迭代，天元拥有了三大核心优势：训练推理一体、全平台高效支持、动静结合的训练能力。

核心优势 1：训练推理一体

所谓训练推理一体，就是让训练与推理不再是完全孤立的两个步骤。为什么训推一体如此重要？在传统深度学习模型开发流程中，从训练到推理，开发者需要经过一系列格式转换。但随之而来的也会有种种复杂难题，比如推理框架对算子不支持、作为中间转换格式的 ONNX 版本不一致，转换成 Caffe 过程中存在一些层的中间形状变化导致无法转换，涉及到量化模型时量化参数的定义不一致等，都会增大转换的难度。

除了转换难外，由于各个平台的差异性和预处理逻辑的复杂性，训练侧提供的预处理逻辑很有可能与推理侧实现的结果不一致，推理和训练的精度对齐变得非常困难。

传统深度学习模型开发环境下，在不断地对分配、适配和微调的过程中，一个模型的落地工作往往需要数天到数周时间。如果还涉及到多平台的部署，则以上问题复杂度和难度会成倍增加，让模型交付变得非常困难。

为了解决传统模型开发的一系列痛点，天元在训练和推理侧均是一套底层框架提供支持，对算子的支持、量化参数的理解均高度一致。同时，天元支持各类常用的 CV 算子操作，便于开发者将预处理的代码逻辑直接写入模型结构当中，从而让训练和推理间的差距大幅度缩小，开发者不再需要为了模型转换而头疼。

经过各方用户一段时间的实际验证，使用天元整个模型，可以将从训练到推理的交付时长缩短至传统方案的十分之一以下，真正做到天级交付。

另外，由于训推一体化，模型在训练阶段与推理阶段可以高度对齐，这使得量化感知训练产生的模型可以无风险的在端侧应用。同时天元拥有精心设计的量化训练模块，大大降低了量化感知训练模型的使用门槛。

以上表格中，在各类模型上，通过应用量化感知训练，可以在极小的精度下降前提下，使用 int8 推理。量化感知训练产生的高精度模型，配合上天元经过高度优化的 int8 算子，可以将各个常见网络的推理耗时降至浮点模型的 30%~50% 以内，有效带来性能加速。

核心优势 2：全平台高效支持

广泛而高效的平台支持，是实现训推一体的前提。只有在各个平台上都提供足够高效便捷的推理能力，才能真正的免除模型转换工作。

对于开发者常用的各类 CPU、GPU 和 NPU，天元均提供了全面而高效的支持，让用户无需为每个平台重新学习重新开发，可以让用户一套模型、一套代码走天下。

天元支持的设备包括 NVIDIA GPU、ARM、X86 和 AMD GPU 平台以及各类国产 NPU 芯片。同时，测试表明，相比其他一些框架，天元在 CUDA、ARM、X86 平台上都具备性能优势。

深度学习应用中，运行时的内存占用是一个重要的指标。由于天元的自动内存优化策略在训练侧得到了充分的打磨，天元可以在内存占用和推理性能之间取得良好的平衡。

仅依靠通用的优化策略，天元对于任意的模型结构都有足够良好的优化表现，框架无需针对网络结构特殊优化。这保证了天元可以在开发者自己的模型结构也保持较低的内存占用，让天元可以帮助各个领域的开发者降低内存占用量，提高竞争力。

端侧推理时，推理程序的二进制体积同样是一项非常重要的指标，对于端侧的 SDK 竞争力非常重要。

对此，天元提供了自动代码裁剪机制，根据模型的结构，将不需要的 kernel 实现和代码片段裁减掉的过程全自动，从而将代码体积大幅度缩小，大幅度提升端侧 SDK 的竞争力。

核心优势 3：动静结合的训练能力

动态图训练对程序员友好，可以随时停下来 Debug；而静态图的训练好处是可以让训练的精度、内存消耗和训练速度达到最优。如何兼得动态图训练和静态图训练的好处？天元通过动静结合支持快速模型设计。动态模式下，天元可以方便的排查模型问题，在遇到错误后及时停下，便于分析模型结构上的异常问题。通过添加一行 trace，模型可以被直接转换为静态图模式，获得静态图所独有的图优化，一方面可以更好的对显存进行复用，提高显存利用率，另一方面可以优化计算顺序，提升模型的训练性能。

在 Beta 版天元推出了亚线性内存优化技术。目前天元是唯一一个支持全自动亚线性内存优化技术的框架。静态模式下，在静态图的基础上，额外添加一个配置项，则可以便捷的启用亚线性内存优化。天元内置的 Sublinear 亚线性内存优化功能，可以对任意的模型结构进行内存优化，在相同的显卡下，可以训练出更大更准确的模型。

亚线性内存优化的特点是参数量越大，从而节省下的显存量越多。因此越庞大的模型，从中获得的收益就越大。由于算法的普适性，无论是任何的模型结构，都可以全自动的被亚线性内存优化机制处理。

共同推进开源生态建设

在使用上，天元模型中心 Model Hub 提供了丰富的预训练模型，包括图像分类、目标检测、图像分割等模型，并且每个模型都提供了 SOTA 级别准确率，让使用者可以便捷的上手天元，站在巨人的肩膀上开始自己的科研工作。

此外，天元已经与小米 MACE、OpenAI Lab Tengine 进行了深度的集成，用户可以将 MegEngine 的模型直接转换到 MACE 或 Tengine 中执行，从而获取在各类异构设备上执行深度学习模型的能力。

深度学习框架的未来路线

天元 1.0 预览版的发布意味着天元已经迭代成为一个完备的深度学习框架。未来天元将如何迭代发展？旷视天元技术团队认为深度学习未来会继续蓬勃发展，在不断涌现新技术新方案的同时，现有的方案和技术会进一步的范式化，简化整个流程，像传统软件开发一样精确分工、快速迭代。因此天元会持续在易用性和整个深度学习落地的流程化上努力，让深度学习应用的流程更加清晰易懂，从而大幅提高 AI 的生产力。

此外，在过去的一段时间中，深度学习框架从编译器中借鉴了相当一部分的内容，旷视天元技术团队认为随着 MLIR 等新兴技术的产生，框架将会进一步向编译器化发展，从而获得更大幅度的性能提升。

最后，随着大量云、端上的 AI 芯片的衍生以及各类神经网络加速器应用的愈加广泛，每家芯片自身的推理框架越来越强大。训推一体的核心是在训练阶段对芯片的充分了解和兼容，因此如何让用户能从一开始就训练出一个适用于最终要用到的芯片的模型，才是让整个落地流程最高效的重点。天元希望能够与芯片协同演进，协同发展，让整个生态能够更加健康蓬勃的发展。

AI 产业落地浪潮中，开发者只有在友好、开放、创新的开发环境中，才能更好地进行 AI 的价值设计与开发，为全球输送 AI 动能。欢迎全球开发者，加入天元社区，一起让深度学习也可以简单开发。

部分摘自于机器之心！

／End.

我们开创“计算机视觉协会”知识星球一年有余，也得到很多同学的认可，我们定时会推送实践型内容与大家分享，在星球里的同学可以随时提问，随时提需求，我们都会及时给予回复及给出对应的答复。

如果想加入我们“计算机视觉研究院”，请扫二维码加入我们。我们会按照你的需求将你拉入对应的学习群！

计算机视觉研究院主要涉及深度学习领域，主要致力于人脸检测、人脸识别，多目标检测、目标跟踪、图像分割等研究方向。研究院接下来会不断分享最新的论文算法新框架，我们这次改革不同点就是，我们要着重”研究“。之后我们会针对相应领域分享实践过程，让大家真正体会摆脱理论的真实场景，培养爱动手编程爱动脑思考的习惯！

计算机视觉研究院

长按扫描二维码
关注我们 获取更多资讯