在近日发布的《开源深度学习框架项目参与指北》文末，我们提到了 MegEngine 在社区开发者的帮助下，已实现了 MegEngine.js —— MegEngine javascript 版本，可以在 javascript 环境中快速部署 MegEngine 模型。
该项目为“开源软件供应链点亮计划-暑期2021”活动项目，本文为 MegEngine.js 项目开发者 - Tricster 所撰写的结项报告的部分节选。enjoy~
方案描述
使用 WebAssembly 将 MegEngine 与 Web 建立联系。
我的实现将保留大部分 C++ 源代码，使用 Typescript 重写 Python 的部分，最后使用 WebAssembly 将 Typescript 和 C++ 连接起来。
这样做的好处是，复用 MegEngine 中的运算符，甚至包括模型的定义和序列化法，可以保证 Megjs 与 MegEngine 最大程度的兼容。
为什么需要 MegEngine.js
造轮子之前，最好先明确这个轮子的价值，避免重复造轮子。而 MegEngine.js 的价值主要体现在两个方面：
端上运算需求大
深度学习不断发展，用户对于自己隐私和数据的保护意识也逐渐增强，如果应用需要将一些敏感数据（身份证照片等）上传到服务器，那用户一定会心有疑虑。边缘设备的计算能力不断增加，也让端上运算变得可行。除了系统层面调用 API 来计算，像微信小程序类需要运行在另一个程序内部，无法直接接触系统 API 的应用，并没有比较合适的方法来计算，许多深度学习应用小程序依然需要将数据发送到服务器上进行计算，在高风险场景下是行不通的。
Web 端需求增大
必须承认的是，Web 有着很强的表达能力，很多新奇的想法都可以在 Web 上进行实现，取得不错的效果，但目前几乎所有的深度学习框架都没有提供 JS 的接口，也就无法运行在 Web 上，如果能比较便捷地在 Web 上运行深度学习框架，将会有很多有趣的应用出现。
MegEngine.js 的架构是什么样的？
想要快速了解一个项目，比较好的方式是先从一个比较 High Level 的角度观察项目的架构、项目中使用的技术，之后再深入代码细节。
大多数深度学习框架的架构
不难发现，几乎所有的深度学习框架其实都有着类似的架构，主要分为三个部分，分别是：
基础运算模块：支持不同设备，不同架构，向上提供统一的接口，高效地完成计算，一般使用C或者C++来编写。
框架主要逻辑模块：在基础运算模块之上，完成深度学习训练和推理的主要逻辑，包括但不限于：计算图的搭建，微分模块的实现，序列化与反序列化的实现，这部分大多也是由C++来编写。
外部接口：由于很多深度学习框架的使用者并不熟悉C++，因此需要在C++之上，创建各种其他语言的绑定，最常见的便是使用Pybind来创建Python绑定。这样一来，使用者便可以在保留Python易用性的情况下，依然拥有良好的性能。
以Pytorch为例，它就是这样的三层结构：
ATen和 C10提供基础运算能力。
由C++实现核心逻辑部分。
将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 上面，用户可以方便地从上面进行下载（https://www.npmjs.com/package/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%。
主要成果展示
demo 1
MegEngine.js Playground，用户可以自由使用 MegEngine.js，测试相关功能（https://codesandbox.io/s/megengine-js-starter-ozf5r）。
demo 2
Megengine.js Model Executor，用户可以加载 MegEngine Model，进行推理。Demo 中所使用的 Model 是通过 MegEngine 官方仓库中示例代码导出的（https://codesandbox.io/s/megengine-js-model-executor-xrbmk）。
demo 3
Megengine.js Linear Regression，线性回归 Demo，展示使用 MegEngine.js 进行动态训练的方式（https://codesandbox.io/s/megengine-js-linear-regression-kur6s）。
demo 4
Megenging.js Mnist，实现了完整的手写数字识别训练与验证（https://codesandbox.io/s/megengine-js-mnist-vhzi8）。
更多 demo
详见仓库中 Example 文件夹（https://gitlab.summer-ospp.ac.cn/summer2021/210040016/-/tree/megjs/megenginejs/example）。
实现 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 的易用性，另外，需要考虑更多边界情况。
点击“阅读原文”，前往知乎阅读更多技术硬核文章
