MegEngine 正式支持 XLA 啦！

XLA（Accelerated Linear Algebra）是 Google 提出的一个神经网络编译器，可以用于加速 AI 模型的训练和推理。MegEngine 1.13.1 中也已经支持了 XLA，在训练模型时可以选择开启此项功能，不同的模型可以获得 10%~80% 不等的速度提升。

主要的目标场景

MegEngine 现在是动态执行的，即 python 中每一个 mge.functional 的调用都对应着底层 gpu 上的一次 kernel 执行。这种模式的好处在于实际的执行方式与代码逻辑一致，所见即所得，非常的灵活；不过其问题是难以优化，性能可能不是最优。

而 XLA 采取静态执行的方式，会将模型计算过程表达成一张静态计算图，称为 “HLO” （High-Level Optimized）。HLO 中包含计算图的相关操作，张量的数据流程和形状等信息。XLA 随后会对 HLO 进行一系列的优化，并最终生成一个更优的计算图，从而更快的完成计算。而 XLA 的局限性就在于不够灵活，对于 Tensor Shape 改变或者控制流等信息无法很好的表达。 

现在 MegEngine 中已经支持了 XLA，模型训练中一些比较静态的场景，我们可以使用 XLA 来进行加速，从而缩短整个训练过程的时间。

使用方法与效果

在使用 MegEngine 进行训练时，可以通过对原来的训练函数增加 xla_trace/partial_trace 装饰器来启用 XLA 编译优化。 

当整个模型是完全静态时，我们可以使用 xla_trace 将整张网络表达成一张静态图，然后交由 XLA 做后续的优化编译，后续的执行过程将执行这张优化后的计算图提升速度。

而如果我们模型中有一些动态性，比如训练过程中一些 Tensor Shape 会发生变化，亦或者是存在控制流，我们可以使用 partial_trace，将网络中静态的部分 trace 成一些子图并分别交给 XLA 进行编译优化，而网络中其他部分仍然保持动态执行，同时保证性能与灵活性。

下面展示了在 MegEngine 中，XLA 功能开启前后，主流的神经网络模型性能变化。其中蓝色为 XLA 开启之前的训练速度，橙色为 XLA 开启之后的训练速度。在开启 XLA 后，大部分模型的性能可以获得 10%~40% 的提升，最多可以超过 80%。

点击 阅读原文 了解关于 XLA 的更多信息及具体的使用方法

M

O

R

E

开源深度学习框架旷视天元（MegEngine）是旷视自主研发的国产工业级深度学习框架，是旷视新一代AI生产力平台Brain++的最核心组件，在 2020 年 3 月正式向全球开发者开源。MegEngine 凭借其训练推理一体、超低硬件门槛和全平台高效推理 3 大核心优势，能够帮助企业与开发者大幅节省产品从实验室原型到工业部署的流程，真正实现小时级的转化能力。

更多 MegEngine 信息获取，您可以：

查看文档：

https://www.megengine.org.cn/doc/stable/zh/

访问深度学习框架 MegEngine 官网：

https://www.megengine.org.cn/

和 GitHub 项目：

https://github.com/MegEngine

同时，也欢迎参与 MegEngine 社区贡献，成为 Awesome MegEngineer，荣誉证书、定制礼品享不停：

https://www.megengine.org.cn/community-AMGE

或加入 MegEngine 用户交流 QQ 群：
1029741705