在日前公布的《开源深度学习框架我的项目参加指北》文末,咱们提到了 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 的角度观察我的项目的架构、我的项目中应用的技术,之后再深刻代码细节。
大多数深度学习框架的架构
不难发现,简直所有的深度学习框架其实都有着相似的架构,次要分为三个局部,别离是:
- 根底运算模块:反对不同设施,不同架构,向上提供对立的接口,高效地实现计算,个别应用
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 下面,用户能够不便地从下面进行下载。
megenginejs
从工作实现状况讲
最后任务书中列出的工作均已实现:
- 能够加载模型和数据
能够间接加载并运行 MegEngine 通过 dump
失去的动态图模型,反对原有框架中的图优化以及存储优化。
-
dense/matmul(必选)的前向 op 单测通过
实现了蕴含
matmul
在内的 21 个常见的 Operator,并全副通过了单元测试。
-
跑通线性回归模型前向,跑通线性回归模型的后向和训练
工作实现,具体实现见
demo3
-
跑通 mnist 模型前向,跑通 mnist 后向和训练
工作实现,具体实现见
demo4
-
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