TensorFlow.js 爲何引入 WASM 後端

在前面的一篇文章《TensorFlow.js 微信小程序插件開始支持 WebAssembly》中,我們談到了 Tensorflow.js(tfjs) 的新後端 WebAssembly(WASM)。這篇文章進一步挖掘 tfjs WASM 後端的更多信息,並探討一下 tfjs 爲何要引入 WASM 後端。

首先澄清一個概念,WASM 後端只是 tfjs 的一種實現,就和 WebGL 實現的後端一樣,對於微信小程序或 WebApp 而言,並不與後端直接打交道,並不會感知到 WASM 後端的存在。

其次,就實現而言,WASM 後端和 WebGL 後端是平行存在的,不存在誰取代誰的問題,針對不同的用戶場景,選擇合適的後端。而實現上的區別, WebGL 後端需要 GPU 支持,且要支持 WebGL ,而 WASM 完全是基於 CPU 運算的。當然,除了這兩個後端,還有一個基於 CPU 運算的普通 Javascript 實現。

我們知道, GPU 非常適合機器學習的運算,不論是模型訓練還是推導, GPU 能從性能上提升好幾個數量級,一般而言機器學習離不開 GPU。那 Google 爲何又要爲 tfjs 開發出一個 WASM 後端呢?這不是在開歷史倒車嗎?

查看了 Google 的官方資料後,總結出如下幾點理由:

  • 大量的低端移動設備缺乏 WebGL 支持,或者有 GPU 但速度很慢。而 WASM 是一種跨瀏覽器工作、可移植彙編和兼容 Web 的二進制文件格式,可在 Web 上實現接近原生代碼的執行速度。全球 90%設備 都支持 WASM。

  • 出於速度上的考慮。普通 Javascript 實現適應面更廣, 但 WASM 的原生解碼速度比 JavaScript 的解析速度 快 20 倍。在前面的文章中也分析過,由於 JavaScript 是動態類型語言,並且會執行垃圾回收,因此可能會在運行時導致明顯的速度問題。而與 WebGL 後端進行對比,雖然對於大多數模型,WebGL 後端的性能仍然優於 WASM 後端,但 WASM 在超精簡模型(小於 3MB 的模型)上的速度則更快。這個時候, GPU 並行化帶來的收益無法抵消執行 WebGL 着色器的固定開銷成本。Google 在 2018 版 MacBook Pro(Intel i7 2.2GHz,Radeon 555X)上,做過 WebGL、WASM 和普通 JS (CPU) 在 Chrome 上的推理時間的對比,結果如下:

從上表可以看出 WASM 後端比普通 JS(CPU)後端快 10-30 倍。與 WebGL 的 PK 則互有勝負,對於類似 MediaPipe 的 BlazeFace 和 FaceMesh 等超輕量級模型,WASM 與 WebGL 速度相當,或比之更快。而對於類似 MobileNet、BodyPix 和 PoseNet 的中型模型,WASM 的速度比 WebGL 慢 2-4 倍。

  • 可移植性和穩定性方面,WASM 原生支持浮點運算,而 WebGL 後端則需要 OES_texture_float 擴展,但並非所有設備都支持此擴展。GPU 驅動程序可能與特定硬件相關,而且不同的設備也可能存在精度問題。比如在 iOS 上,GPU 不支持 32 位浮點數,因此只能退化到 16 位浮點數,從而導致精度問題。在 WASM 中,將始終以 32 位浮點數進行計算,因此在所有設備之間都能實現一致的精度。

  • WASM 仍然存在優化的空間。值得一體的方案有 SIMD(Single Instruction Multiple Data,單指令流多數據流) WASM 擴展方案。在熱門機器學習模型上使用 SIMD-WASM 進行的基準測試表明,其速度相比非 SIMD WASM 提高了 2-3 倍,而採用 LLVM 優化 SIMD 指令後,還可額外提速 26-50%。另外多線程的支持更加完善後,可充分利用線程提升性能,而無需修改 tfjs 用戶代碼。

小結

tfjs WASM 後端因其設備支持廣泛,在性能上又能取得不錯的平衡,特別是對於微信小程序而言,通常不會選用太複雜的模型,性能優勢更明顯。而 SIMD 和線程之類的新擴展,將如虎添翼,讓 tfjs WASM 後端越來越受歡迎。作爲一名 C++ 程序員,我也希望我的 C++ 編程技能也能在 Web 應用開發方面一展拳腳。

參考

  1. Introducing the WebAssembly backend for TensorFlow.js

  2. Face and hand tracking in the browser with MediaPipe and TensorFlow.js

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章