在日前釋出的《開源深度學習框架專案參與指北》文末，我們提到了 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