本文整顿自 2 月 25 号 MegEngine 首期 Meetup 的分享。作者：王彪 | 旷视框架部异构计算组工程师视频回放：《JIT in MegEngine》_MegEngine Meetup 第一期一、背景什么是天元旷视天元（MegEngine）是一个深度学习框架，它次要蕴含训练和推理两方面内容。训练侧个别应用 Python 搭建网络；而推理侧思考到产品性能的因素，个别应用 C++ 语言集成天元框架。无论在训练侧还是推理侧，天元都负担着将训练和推理的代码运行到各种计算后端上的工作。目前天元反对的计算后端有 CPU、GPU、ARM 和一些畛域专用的加速器，笼罩了云、端、芯等各个场景。
天元次要有三大特色：
1.训推一体，不论是训练任务还是推理工作都能够由天元一个框架来实现。
2.动静联合，天元同时反对动态图和动态图，并且动静之间的转换也十分不便。
3.多平台的高性能反对。 如图 1 所示，咱们能够看到天元提供了 Python 和 C++ 两种接口。在图示意上分为动态图和动态图。运算层组件包含主动求导器、图优化和图编译等。天元的运行时模块包含内存治理和计算调度，其中内存治理包含动态内存治理和动态内存治理，以及亚线性内存优化技术。计算内核层蕴含了天元反对的所有计算后端，咱们后续会开源出更多的计算后端。除此之外，天元还蕴含了一个高性能异构通信库，它个别会在多机多卡的场景下被用到。
动态图和动态图是绝对的，在动态图下是没有计算图的概念的。但在动态图下，天元会保护一张计算图。如图 2 所示为天元中的计算图示意，图中圆形示意算子（operator），三角形示意输出。在天元框架中，动态图和动态图之间的转换只须要一条简略的语句即可实现，如下方代码所示：
if __name__ == '__main__’: gm = ad.GradManager().attach(model.parameters()) opt = optim.SGD(model.parameters(), lr=0.0125, momentum=0.9, weight_decay=1e-4) <em> # 通过 trace 转换为动态图</em> @trace(symbolic=True) def train(): with gm: logits = model(image) loss = F.loss.cross_entropy(logits, label) gm.backward(loss) opt.step() opt.clear_grad() return loss loss = train() loss.numpy()什么是 AOT 和 JITAOT(Ahead Of Time) 和 JIT(Just In Time) 都是编译中的概念。以传统的 C/C++ 语言为例，咱们写完代码之后，个别会通过编译器编译生成可执行文件，而后再执行该可执行文件取得执行后果。如果咱们将从源代码编译生成可执行文件的过程称为 build 阶段，将执行可执行文件叫做 runtime 阶段的话，JIT 是没有build 阶段的，它只有 runtime 阶段。JIT 个别被用在解释执行的语言如 Python 中，JIT 会在代码执行的过程中检测热点函数，随后对热点函数进行重编译，下次运行时遇到热点函数则间接执行编译后果即可。这样做能够显著放慢代码执行的速度。
...