Web 上运行深度学习框架模型 - MegEngine.js

在日前发布的《开源深度学习框架项目参与指北》文末，我们提到了 MegEngine 在社区开发者的帮助下，已实现了 MegEngine.js —— MegEngine javascript 版本，可以在 javascript 环境中快速部署 MegEngine 模型。

该项目为“开源软件供应链点亮计划 - 暑期2021”活动项目，本文为 MegEngine.js 项目开发者 - Tricster 所撰写的结项报告的部分节选。enjoy~

项目信息

方案描述

使用 WebAssembly 将 MegEngine 与 Web 建立联系。

我的实现将保留大部分 C++ 源代码，使用 Typescript 重写 Python 的部分，最后使用 WebAssembly 将 Typescript 和 C++ 连接起来。

这样做的好处是，复用 MegEngine 中的运算符，甚至包括模型的定义和序列化方法，可以保证 MegEngine.js 与 MegEngine最 大程度的兼容。

为什么需要 Megengine.js ？

造轮子之前，最好先明确这个轮子的价值，避免重复造轮子。而 Megengine.js 的价值主要体现在两个方面：

端上运算需求增大

深度学习不断发展，用户对于自己 隐私 和 数据 的保护意识也逐渐增强，如果应用需要将一些敏感数据（身份证照片等）上传到服务器，那用户一定会心有疑虑。边缘设备的计算能力不断增加，也让端上运算变得可行。除了系统层面调用 API 来计算，像微信小程序这类需要运行在另一个程序内部，无法直接接触系统 API 的应用，并没有比较合适的方法来计算，许多深度学习应用小程序依然需要将数据发送到服务器上进行计算，在高风险场景下是行不通的。

Web 端需求增大

必须承认的是，Web 有着很强的表达能力，很多新奇的想法都可以在 Web 上进行实现，取得不错的效果，但目前几乎所有的深度学习框架都没有提供 JS 的接口，也就无法运行在 Web 上，如果能比较便捷地在 Web 上运行深度学习框架，将会有很多有趣的应用出现。

Megengine.js 的架构是什么样的？

想要快速了解一个项目，比较好的方式是先从一个比较 High Level 的角度观察项目的架构、项目中使用的技术，之后再深入代码细节。

大多数深度学习框架的架构

不难发现，几乎所有的深度学习框架其实都有着类似的架构，主要分为三个部分，分别是：

1. 基础运算模块：支持不同设备，不同架构，向上提供统一的接口，高效地完成计算，一般使用 C 或者 C++ 来编写。
2. 框架主要逻辑模块：在基础运算模块之上，完成深度学习 训练 和 推理 的主要逻辑，包括但不限于： 计算图的搭建 ， 微分模块的实现 ， 序列化与反序列化的实现 ，这部分大多也是由 C++ 来编写。
3. 外部接口：由于很多深度学习框架的使用者并不熟悉 C++，因此需要在 C++ 之上，创建各种其他语言的 绑定 ，最常见的便是使用 Pybind 来创建 Python 绑定。这样一来，使用者便可以在保留 Python 易用性的情况下，依然拥有良好的性能。

以 Pytorch 为例，它就是这样的三层结构：

1. ATen 和 C10 提供基础运算能力。
2. 由 C++ 实现核心逻辑部分。
3. 将 C++ 部分作为 Extension 供 Python 调用，只在 Python 中进行简单的包装。

以 MegEngine 为例

MegEngine 文件结构还是比较清晰的，主要如下：

.
├── dnn
├── imperative
└── src 

虽然 MegEngine 有类似的结构，但依然有些不同。

dnn文件夹中的MegDnn，是底层运算模块，支持不同架构、不同平台，比如x86、CUDA、arm。这些模块虽然实现方式各不相同，但是都提供了统一的接口，供 MegEngine 调用。

如右图所示，不同架构的算子按包含关系组成了一个树形结构。虽然现在一般都是使用叶节点的算子，但naive和fallback 在开发过程中也是相当重要的部分，对实现新的算子有很大的帮助。

另外，采样这样的树形结构的，可以很好地复用代码，比如我们可以只实现部分算子，其他算子可以向上寻找已有的实现，可以节省很多的工作量。

MegDnn 算子组织架构图

src中包含了 MegEngine 的主要代码，核心是如何构建一个计算图（静态图）以及Tensor的基础定义，除此之外，还对存储，计算图进行很多的优化，简单来说，只用MegDnn以及src中的代码，可以进行高效地运算（Inference Only），并不包含训练模型所需要的部分，更多地用于部署相关的场景。

最后imperative中，补全了一个神经网络框架的其他部分，比如反向传播、各种层的定义以及一些优化器，并且使用Python向外提供了一个易用的接口。值得一提的是，在imperative中，使用 Pybind 将 C++ 和 Python 进行了深度耦合， Python 不再只作为暴露出来的接口，而是作为框架的一部分，参与编写了执行逻辑。比如动态 计算图转换成静态计算图 这个功能，就是一个很好的例子，既利用了Python中的装饰器，又与C++中静态计算图的部分相互配合。

采用这样的架构，是比较直观且灵活的，如果想要增加底层运算模块的能力，只需要修改MegDnn就好；如果想增加静态图相关的特性，只需要修改src的部分；如果想要对外添加更多的接口功能，只修改imperative便可以做到。

理论上来讲，如果想要将 MegEngine Port 到其他语言，只需要替换掉imperative就可以，但由于imperative中C++与Python耦合比较紧密，就必须先剥离所有 Python 的部分，然后再根据需要补上目标语言的实现（C++、JS 或是其他语言）。

MegEngine.js 设计思路

基于上述的分析，Megengine.js 采用了下图的架构。

底层复用 MegEngine 的实现，包括计算模块，以及计算图的实现；然后模仿 Python 的部分使用 C++编写一个 Runtime ，完成 imperative中提供的功能，并存储所有的状态；然后使用 WebAssembly 将上述所有模块暴露给 TypeScript 来使用，并用 TypeScript实现剩余的逻辑部分，提供一个易用的接口给用户来使用。

采用这样的架构，最大程度将 MegEngine.js 作为一个顶层模块融入 MegEngine ，而不是像 Tensorflow.js 那样从头实现一个 Web 端的深度学习框架。这样做的好处是， MegEngine.js 不仅可以享受到 MegEngine 高度优化之后的特性，还可以直接运行 MegEngine 训练的模型，为之后的部署也铺平了道路。

Megengine.js 现在处于什么状况？

从框架角度讲

目前 MegEngine.js 已经是一个可以正常使用的框架了，验证了整个实现方案的可行性。用户可以使用 MegEngine.js 直接运行从 MegEngine 导出的静态图模型，也可以从头搭建一个自己的网络，在浏览器中进行训练，推理，并且可以加载和保存自己的模型，除此之外，用户也可以在 Node 的环境中进行上述任务。

MegEngine.js 已经发布到 NPM 上面，用户可以方便地从上面进行下载。

megenginejs

从任务完成情况讲

最初任务书中列出的任务均已完成 ：

1. 可以加载模型和数据

可以直接加载并运行 MegEngine 经过 dump 得到的静态图模型，支持原有框架中的图优化以及存储优化。

2. dense/matmul （必选）的前向 op 单测通过

   实现了包含 matmul 在内的21个常见的 Operator ，并全部通过了单元测试。

3. 跑通线性回归模型前向，跑通线性回归模型的后向和训练

   任务完成，具体实现见 demo3

4. 跑通 mnist 模型前向，跑通 mnist 后向和训练

   任务完成，具体实现见 demo4

5. mnist 的 demo

   完成了 mnist 的训练以及验证，但并未实现相关可视化（损失变化，准确率变化，测试样本），见 demo4

解决性能瓶颈

除此之外，由于 WebAssembly的限制和 Web 跨平台的特性，我无法使用 MegEngine 中高度优化的算子，导致在初期性能表现并不理想，无法带来流畅的体验，于是在中期之后，我参考 Tensorflow.js，引入了 XNNPACK ，实现了一套新的算子，有效地提升了 Megengine.js 的运行速度。

在 MacOS 平台进行算子的 Benchmark ，卷积算子的运行耗时降低 83% 。

WASM.BENCHMARK_CONVOLUTION_MATRIX_MUL (6169 ms)
WASM.BENCHMARK_CONVOLUTION_MATRIX_MUL (36430 ms)

在 Safari 中进行 Mnist 训练，单次训练时间下降 52% 。

主要成果展示

Demo1

Megengine.js Playground ，用户可以自由使用 Megengine.js ，测试相关功能。

Megengine.js Starter

Demo2

Megengine.js Model Executor ，用户可以加载 MegEngine Model ，进行推理。Demo中所使用的Model是通过 MegEngine 官方仓库中示例代码导出的。

Megengine.js Model Executor

Demo3

Megengine.js Linear Regression ，线性回归 Demo ，展示使用 MegEngine.js 进行动态训练的方式。

Megengine.js Linear Regression

Demo4

Megenging.js Mnist ，实现了完整的手写数字识别训练与验证。

Megengine.js Mnist

更多Demo

详见仓库中 Example 文件夹。

megenginejs/example · megjs · Summer2021 / 210040016

实现 Megengine.js 的过程中遇到了什么样的问题？

虽然从一开始就设想好了架构，分层也比较明确，但仍然遇到了许多问题。

编译问题

问题描述

MegEngine 是使用 C++ 编写的，所以第一步就应该是将 MegEngine 编译为 WebAssembly ，借助 Emscripten 可以将简单的 C++ 程序编译成 WASM，但对于 MegEngine 这样体量的项目，就没办法不更改直接编译了。

解决办法

最大的问题，主要是 MegDnn 这个算子库包含了太多平台依赖的部分和优化，在尝试很多方案后，还是没有办法将那些优化也包含进来，于是最后只能先去掉所有的优化，使用最直接的实现方式（Naive Arch），关掉其他一些编译选项之后，完成了编译。

但在这里的处理不得已选择了 速度比较慢 的算子，也导致框架的整体速度不太理想。

交互问题

问题描述

无论是 MegEngine 还是 Megengine.js ，都需要让 C++ 编写的底层与其他语言来进行交互。使用 Pybind的时候，可以比较紧密地将 C++ 和 Python 结合起来，在 Python 中创建、管理 C++ 对象，但在 Emscripten 这边，要么使用比较底层的 ccall 和 cwrap ，要么使用 Embind 来将 C++ 对象与 Python 进行绑定， Embind虽然模仿 Pybind，但没有提供比较好的 C++ 对象的管理方法，所以没办法像 Pybind 那样把 Python 和 C++ 紧耦合起来。

最理想的情况下，JS 与C++应该管理同一个变量，比如Python创建的Tensor，继承了C++的Tensor，当一个Tensor在Python中退出作用域，被GC回收时，也会直接销毁在C++中创建的资源。这样的好处也相当明显，Tensor可以直接作为参数在C++与Python之间来回传递，耦合很紧密，也非常直观。

但是在JS中，这是做不到的，首先cwrap和ccall只支持基本类型，Embind虽然支持绑定自定义的类，但是使用起来比较繁琐，用这种方法声明的变量还必须手动删除，增加了许多负担。

解决办法

在这种情况下，我选择在 C++ 里面内置一个 Runtime ，用这个 Runtime 来管理 Tensor 的生命周期，并且用来追踪程序运行中产生的状态变量。

比如在JS中创建Tensor后，会将实际的数据拷贝到C++中，在C++创建实际管理数据的 Tensor （也是 MegEngine 中使用的Tensor），之后交给 C++ Runtime 来管理这个 Tensor ，创建好后，将这个 Tensor 的 ID 返回给 JS 。也就是说， JS 中的 Tenosr 更像是一个指针，指向 C++ 中的那个 Tensor 。

这样进行分离后，虽然需要管理 C++ 和 JS 中 Tensor 的对应关系，但这样大大简化了 JS 和 C++ 之间的调用，无论是使用基础的 ccall 、 cwrap 还是 Embind 都可以传递 Tensor 。

当然，这样做也有弊端，由于 C++ 和 JS 是分离的设计，需要写不少重复的函数。

GC问题

问题描述

JS 和 Python 都是有 GC 的， Python 在 MegEngine 中发挥了很大的作用，可以及时回收不再使用的 Tensor ，效率比较高，但是在 JS 中的情况更加复杂。虽然 JS 是有 GC 的，但是与 Python 激进的回收策略相比， JS\ 更加佛系，可能由于\ 浏览器 的使用场景或是 JS 的设计哲学。一个变量是否被回收，何时被回收，都没有办法被确定，甚至在一个变量被回收的时候，都没有办法执行一个回调函数。

解决办法

为了解决这个问题，我只能实现一个朴素的标记方法，将跳出 Scope 的变量回收掉，避免在运行过程中内存不够用的情况。但这种朴素的方法还是有些过于简单了，虽然确实可以避免内存溢出的情况，但依然效率不算高。

关于Finalizer

在 JS 新的标准中，增加了一个机制，可以让我们在一个变量被 GC 回收时调用一个回调函数（ Finalizer ），来处理一些资源。理想很美好，实际测试中，这个变量被回收的时间是很不确定的（ JS 的回收策略比较佛系），不仅仅如此，我们的 Tensor 数据实际存储在 WebAssembly 之中的， JS 的 GC 并不能监控 WASM 中的内存使用情况，也就是说即使 WASM 中内存被占满了，由于 JS 这边内存占用还比较少， GC 并不会进行回收。

基于这两点原因， Finalizer 并不是一个很好的选择。

P.S. 很多浏览器还不支持 Finalizer。

性能问题

问题描述

之前提到，为了成功将 MegEngine 编译成 WebAssembly ，牺牲了很多东西，其中就包含高性能的算子，虽然整个框架是可以运行的，但是，这个效率确实不能满足用户的正常使用。问题的原因很简单， MegEngine 中并没有针对 Web 平台进行的优化，所以为了解决这个问题，只能考虑自己实现一套为 Web 实现的算子。

解决方法

专为 Web 进行优化的 BLAS 其实不算多， Google 推出的 XNNPACK 是基于之前 Pytorch 推出的 QNNPACK 上优化的，也被用在的 Tensorflow.js 里面，所以我这里选择将 XNNPACK 加入进来。但由于 XNNPACK 里面的诸多限制，并没有加入全部的算子，但改进之后速度还是有了不错的提升。

Megengine.js 之后会如何发展？

经过 3 个月的开发，对 MegEngine 的了解也越发深入，也越来越想参与到社区的建设中来。 Megengine.js 虽然具备了基础的功能，但距离一个完整的框架还有不小的差距，之后还有许多工作可以做。

进一步完善各种模块

一个合格的深度学习框架应该有比较全面的算子支持，模块支持，现在 MegEngine.js 支持的算子和模块还是比较少的，之后还需要再添加更多的实用的算子，这样才能利于这个框架的进一步推广。

进一步提升性能

对性能的提升是永远不够的，在这样一个浮躁的时代，运行速度是一个不可忽视的指标。虽然 XNNPACK 的加入提升了速度，但其实还不够，不仅仅是因为算子的支持不够，而且应该还是有更多的提升空间的。

进一步优化框架

不要过度优化，但是也不能让代码变成一潭死水，在合适的时候（完成必要的功能模块之后），可能需要进一步提升 Megengine.js 的易用性，另外，需要考虑更多边界情况。

延展阅读

【作者博客】Web 上的深度学习 | Avalon

【教程】小程序中使用 MegEngine.js 教程

欢迎更多的开发者加入到 MegEngine 社区，这里还有一份适合新手的参与教程及任务清单：

开源深度学习框架项目参与指北 - 内含易上手任务清单

MegEngine 技术交流群，QQ 群号：1029741705