MegEngine 动态执行引擎 Imperative Runtime 架构解析

最新推荐文章于 2025-12-05 13:34:07 发布

原创

最新推荐文章于 2025-12-05 13:34:07 发布 · 941 阅读

19 ·

CC 4.0 BY-SA版权

文章标签：

#开源 #深度学习 #python

本文详细介绍了MegEngine的Python层模块，包括data、functional、module、optimizer、quantization和tools，以及ImperativeRuntime中关键组件的工作原理，如dispatcher的预处理和interpreter的计算执行。此外，还探讨了dispatcher如何处理transformation，以及Interpreter在执行op时的角色和Transformer间的执行顺序。

在之前的文章中我们介绍过 MegEngine 的 Imperative Runtime 以及它与 MegBrain、MegDNN 的关系，这篇文章中我们将介绍 Imperative 中包含的常用组件。

在 MegEngine 中，从用户在 python 层编写代码到在 interpreter 层发生计算经过了下面的流程：

用户在 python 层编写网络结构代码，执行时向 C++ 层发射算子执行指令
Imperative 的 dispatcher 对部分算子做计算前的预处理（transformation）
Imperative 的 interpreter 执行计算操作（复用 MegBrain 的相关组件）

我们将分别介绍这几个阶段系统所做的工作。

MegEngine 的 Python 层

在主流的深度学习框架中，用户往往不需要自己手写算子的具体实现、处理计算图的执行逻辑、或者与复杂的体系结构打交道。一切都被封装为 Python 层的接口。

在 MegEngine 的 Python 层中用户接触较多的模块主要有：data、functional、module、optimizer、quantization、tools，下面简单介绍一下各个模块的功能。

构建数据处理 Pipeline —— data 模块

Data 模块，顾名思义就是对数据进行处理的模块。

没有数据就没法训练，在 MegEngine 中，通常会借助一个 Dataset 结构来定义数据集。数据集一般分为 Map-stype 和 Iterable-style 两种。前者叫作 ArrayDataset，这种数据集支持随机访问；后者叫作 StreamDataset，因为是流式的数据集，只支持顺序访问。

有了数据集，我们还需要一个结构来把数据“喂”给模型训练，这样的一个结构叫作 dataloader。

实际上，只给 dataloader 一个 dataset 有时无法准确地描述加载数据的整个过程，我们可能还需要定义加载数据过程的抽样规则（