在前几天发布的《开源深度学习框架项目参与指北》文章的最后，我们提到MegEngine在社区开发者的帮助下，已经实现了MegEngine.js——MegEnginejavascript版本，可以在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上实现，并取得了不错的效果。但是目前几乎所有的深度学习框架都不提供JS接口。它不能在Web上运行。如果深度学习框架能够更方便地在Web上运行，将会有很多有趣的应用。Megengine.js的架构是怎样的？如果想快速了解一个项目，更好的方法是先从高层次的角度去观察项目的架构和项目中使用的技术，然后再深入代码细节。大多数深度学习框架的架构都不难找到。几乎所有的深度学习框架其实都有相似的架构，主要分为三个部分，即：基础计算模块：支持不同设备，不同架构，向上提供统一接口，高效完成计算，一般用C或C语言编写C++。框架主要逻辑模块：在基础计算模块之上，完成深度学习训练和推理的主要逻辑，包括但不限于：计算图的构建、微分模块的实现、实现序列化和反序列化，这部分大部分也是用C++写的。外部接口：由于很多深度学习框架的用户不熟悉C++，他们需要在C++之上为各种其他语言创建绑定。最常见的一种是使用Pybind创建Python绑定。这样一来，用户在保留Python的易用性的同时，仍然可以拥有良好的性能。以Pytorch为例，就是这样一个三层结构：ATen和C10提供基础的计算能力。核心逻辑部分由C++实现。将C++部分作为Python的Extension来调用，只在Python中进行简单的封装。以MegEngine为例，MegEngine的文件结构比较清晰，主要有：.├──dnn├──imperative└──src虽然MegEngine的结构类似，但还是有一些区别。dnn文件夹下的MegDnn是底层计算模块，支持x86、CUDA、arm等不同架构和平台。这些模块虽然实现方式不同，但都提供了统一的接口供MegEngine调用。如右图所示，不同架构的算子按照包含关系形成树状结构。虽然现在普遍使用叶子节点算子，但是naive和fallback也是开发过程中非常重要的部分，对实现新的算子有很大的帮助。另外，采样这样的树结构可以很好的重用代码。比如我们可以只实现一些算子，其他算子可以向上搜索已有的实现，这样可以节省很多工作量。MegDnn运营商组织结构图src包含了MegEngine的主要代码。核心是如何构建计算图（静态图）和Tensor的基本定义。此外，它还优化了存储和计算图。也就是说，只有MegDnn和src中的代码可以用于高效计算（InferenceOnly），不包括训练模型所需的部分，更多用于部署相关的场景。最后在命令式中完成了一个神经网络框架的其他部分，比如反向传播，各层的定义，以及一些优化器，并使用Python对外提供了一个简单易用的接口。值得一提的是，在命令式中，C++和Python使用Pybind深度耦合。Python不再只是一个暴露的接口，而是参与编写执行逻辑的框架的一部分。比如将动态计算图转换为静态计算图的功能就是一个很好的例子，它不仅使用了Python中的装饰器，还配合了C++中的静态计算图。采用这样的架构是比较直观和灵活的。如果想增加底层计算模块的能力，只需要修改MegDnn即可；如果要添加静态图片相关的功能，只需要修改src部分即可；如果想增加更多的接口功能，只需要修改命令式就可以实现。从理论上讲，如果要将MegEngine移植到其他语言，只需要替换命令式即可，但由于C++和Python在命令式上紧耦合，因此必须先剥离所有Python部分，然后根据需要添加目标语言实现（C++、JS或其他语言）。MegEngine.js设计思路基于以上分析，Megengine.js采用了如下图所示的架构。底层复用了MegEngine的实现，包括计算模块和计算图的实现；然后模仿Python部分，用C++写一个Runtime，完成命令式提供的功能，存储所有的状态；然后使用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的任务完成情况来看，初始任务书列出的任务已经完成：可以加载模型和数据，转储MegEngine得到的静态图模型可以直接加载运行，支持图优化和存储在原有框架下进行优化。dense/matmul(required)的forwardop单元测试通过matmul等21个常用Operator的实现，全部通过单元测试。正向运行线性回归模型，反向运行线性回归模型，完成训练任务。具体实现见demo3。mnist训练和验证，但是没有实现相关的可视化（loss变化，accuracy变化，测试样本），看demo4解决性能瓶颈另外，由于WebAssembly的限制和Web的跨平台特性，我无法使用MegEngine高度优化的算子导致初期性能不尽如人意，无法带来流畅的体验。所以在中期之后，我参考了Tensorflow.js引入了XNNPACK，实现了一套新的算子，有效的提升了Megengine.js的速度。在MacOS平台上算子的Benchmark中，卷积算子的运行时间减少了83%。WASM.BENCHMARK_CONVOLUTION_MATRIX_MUL(6169ms)WASM.BENCHMARK_CONVOLUTION_MATRIX_MUL(36430ms)Mnist在Safari中训练，单次训练时间下降52%。主要成果展示Demo1Megengine.jsPlayground，用户可以自由使用Megengine.js测试相关功能。Megengine.jsStarterDemo2Megengine.jsModelExecutor，用户可以加载MegEngine模型进行推理。Demo中使用的Model是从MegEngine官方仓库的示例代码中导出的。Megengine.jsModelExecutorDemo3Megengine.jsLinearRegression线性回归Demo，展示了如何使用MegEngine.js进行动态训练。Megengine.jsLinearRegressionDemo4Megenging.jsMnist实现了完整的手写数字识别训练和验证。有关Megengine.jsMnist的更多演示，请参阅存储库中的示例文件夹。megenginejs/example·megjs·Summer2021/210040016在实现Megengine.js的过程中遇到了什么样的问题？虽然从一开始就构思了架构，层次也比较清晰，但是还是遇到了很多问题。编译问题问题描述MegEngine是用C++写的，所以第一步应该是把MegEngine编译成WebAssembly。借助Emscripten，可以将一个简单的C++程序编译成WASM，但是对于MegEngine这种规模的项目，没办法不改。直接编译。解决方案中最大的问题是MegDnn算子库包含过多的平台相关部分和优化。尝试了很多方案，还是没有办法把那些优化都包含进去，所以最后只能先去掉所有优化。使用最直接的实现方式（NaiveArch），关闭其他一些编译选项后，编译完成。但是这里的处理不得不选择一个比较慢的算子，这也导致了框架整体速度不尽如人意。交互问题问题描述无论是MegEngine还是Megengine.js，都需要让C++编写的底层能够与其他语言进行交互。使用Pybind时，可以更紧密地结合C++和Python，在Python中创建和管理C++对象，但在Emscripten上，要么使用更底层的ccall和cwrap，要么使用Embind将C++对象与Python结合起来。Binding，Embind虽然模仿了Pybind，但是并没有提供更好的C++对象的管理方式，所以没有办法像Pybind那样将Python和C++紧密耦合。理想情况下，JS和C++应该管理同一个变量。比如Python创建的Tensor继承了C++的Tensor。在Python中一个Tensor退出作用域，被GC回收时，也会直接销毁。CreatedinC++H。这样做的好处也是相当明显的。Tensor可以直接作为参数在C++和Python之间来回传递。耦合非常紧密且非常直观。但是在JS中，这是不可能的。首先，cwrap和ccall只支持基本类型。Embind虽然支持绑定自定义类，但是使用起来比较麻烦。这种方式声明的变量必须手动删除，增加了很多负担。解决方案在这种情况下，我选择用C++构建一个Runtime，用这个Runtime来管理Tensor的生命周期，用它来跟踪程序运行过程中产生的状态变量。比如在JS中创建一个Tensor后，实际的数据会被复制到C++中，而实际管理数据的Tensor（也是MegEngine中使用的Tensor）会在C++中创建，然后交给C++Runtime来实现管理张量。创建完成后，将这个Tensor的ID返回给JS。也就是说，JS中的Tenosr更像是一个指向C++中的Tensor的指针。这样分离之后，虽然需要管理C++中的Tensor和JS的对应关系，但是这大大简化了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中的内存满了，GC也不会回收，因为JS上的内存使用方还是比较小的。由于这两个原因，Finalizer不是一个好的选择。附言许多浏览器还不支持Finalizer。Performanceproblem问题描述前面提到，为了将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的可用性。此外，还需要考虑更多的边界条件。延伸阅读【作者博客】DeepLearningontheWeb|Avalon【教程】MegEngine.js小程序使用教程欢迎更多开发者加入MegEngine社区，这里还有适合新手的参与教程和任务列表：开源深度学习框架项目参与指南-包含易于上手的列表-使用任务MegEngine技术交流群，QQ群号：1029741705
