如何打造一個高性能的前端智能推理引擎

什麼是前端智能推理引擎

在前端智能推理引擎之前,我們先來說一下什麼是”端智能”。

端智能(On-Device Machine Learning)是指把機器學習的應用放在端側做。這裏的“端側”,是相對於雲服務而言的。它可以是手機,也可以是 IOT 設備等。

傳統的機器學習,由於模型大小、機器算力的問題,很多是放在服務端做的。比如 Amazon AWS 有“Amazon Rekognition Service”,Google 有 “Google Cloud Vision Service”。而隨着以手機爲代表的端側設備算力的提高,以及模型設計本身的演進,大小更小、能力更強的模型逐漸能夠部署到端上運行。

相比雲端部署的方式,APP端擁有更直接的用戶特徵,同時具備如下優勢:

  • 實時性高,端側處理可節省數據的網絡傳輸時間。
  • 節省資源,充分利用端側算力和存儲空間。
  • 隱私性好,產生數據到消費數據都在端側完成,避免傳輸引起的隱私泄露風險。

這些是端智能的優勢,但它不是萬金油,仍然存在一些侷限性:

  • 設備資源有限,端側算力、存儲是有限的,不能做大規模高強度的持續計算。
  • 算法規模小,端側算力小,而且單用戶的數據,在算法上並不能做到最優。
  • 用戶數據有限,端側數據不適合長期存儲,同時可用數據有限。

同理,前端智能是指將機器學習的應用放到前端上(web、h5、小程序等).

所以,什麼是前端智能推理引擎呢?

如下圖:

前端智能推理引擎實際上就是利用前端上算力去執行模型的那個東西。

業界現有的前端推理引擎

這裏列出三個常見的推理引擎

  • tensorflow.js(下面簡稱爲tfjs)
  • ONNX.js
  • WebDNN

對於一個端上推理引擎來說,最重要的是什麼?當然是性能了!性能越好,也代表在端上的應用場景也會越多,下面我們來看下這三個推理引擎的性能對比:

(下面數據使用模型爲MobileNetV2分類模型)

cpu(js計算)

可以看到,在純JS環境下進行計算,僅僅做一次分類都要1500ms以上。設想一下如果一個相機需要實時對拍攝的物體做分類預測(比如預測拍攝的對象是貓還是狗),那麼每預測一次需要1500ms,這樣的性能是無法忍受的。

WASM

在WASM環境下,性能最佳的ONNX.js達到了135ms的性能,也就是7fps左右,已經到了勉強能用的程度了。而tfjs卻是糟糕的1501ms。這裏是因爲onnx.js利用了worker進行多線程加速,所以性能最好。

WebGL(GPU)

最後是GPU環境,可以看到tfjs和ONNXjs的性能都達到了比較好的性能水平,而WebDNN表現較爲糟糕。

除了上面這三種引擎,目前國內還有百度的paddle.js以及淘寶的mnn.js等,這裏不做討論。

當然,在選擇一個合適的推理引擎時,除了性能以外,還有生態、引擎維護情況等等一系列的考慮。從綜合的方面來說,tfjs是當下市場上最適合的前端推理引擎。因爲tfjs可以依靠tensorflow的強大的生態、google官方團隊的全職維護等。相比之下ONNX框架比較小衆,且ONNXjs已經有近一年沒有維護了。WebDNN性能及生態都沒有任何競爭力。

前端上的高性能計算方案

從上一章節其實能看到,在前端上做高性能計算一般比較普遍的就是WASM和基於WebGL的GPU計算,當然也有asm.js這裏不做討論。

WASM

WASM大家應該是比較熟悉的,這裏只做下簡短的介紹:

WebAssembly是一種運行在現代網絡瀏覽器中的新型代碼,並且提供新的性能特性和效果。它設計的目的不是爲了手寫代碼而是爲諸如C、C++和Rust等低級源語言提供一個高效的編譯目標。

對於網絡平臺而言,這具有巨大的意義——這爲客戶端app提供了一種在網絡平臺以接近本地速度的方式運行多種語言編寫的代碼的方式;在這之前,客戶端app是不可能做到的。

而且,你在不知道如何編寫WebAssembly代碼的情況下就可以使用它。WebAssembly的模塊可以被導入的到一個網絡app(或Node.js)中,並且暴露出供JavaScript使用的WebAssembly函數。JavaScript框架不但可以使用WebAssembly獲得巨大性能優勢和新特性,而且還能使得各種功能保持對網絡開發者的易用性。--《摘自MDNWebAssembly概念》

WebGL

啥?WebGL不是做圖形渲染的嗎?不是做3D的嗎?爲啥能做高性能計算?

可能一些同學聽說過gpgpu.js這個庫,這個庫就是利用webgl做通用計算的,具體的原理是怎麼樣的呢?(爲了能夠繼續往下閱讀,請先快速瀏覽下這篇文章):《利用WebGL2 實現Web前端的GPU計算》。

將推理引擎的性能進行極致優化

好了,目前我們知道在前端上的兩種高性能計算方式了,那麼如果現有的框架(tfjs、onnxjs)性能上就是不滿足我們的需求怎麼辦呢?怎麼樣才能進一步提升引擎性能,並落地生產環境呢?

答案是:手撕源碼,優化性能。對,就是這麼簡單粗暴。以tfjs爲例(其他的框架原理上是一致的),下面給大家介紹下如何用不同的姿勢去優化引擎性能。

在去年年初時候,我們團隊和google的tfjs團隊做了一次深入交流,google那邊明確表示tfjs後面的發展方向以WASM計算爲主、webgl計算不做新的feature以維護爲主。但是現階段各瀏覽器、小程序對WASM的支持並不完整(例如SIMD、Multi-Thread等特性),所以WASM暫時無法在生產環境落地。所以,現階段還是需要依賴webgl的計算能力。糟糕的是,此時tfjs的webgl性能在移動端上表現依舊差強人意,尤其在中低端機上的性能完全達不到我們的業務要求。沒辦法,只能自己硬着頭皮進去優化引擎。所以以下的內容都是針對於webgl計算進行介紹。

優化 WebGL 高性能計算的n種姿勢

姿勢一:計算向量化

計算向量化是指,利用glsl的vec2/vec4/matrix數據類型進行計算,因爲對於GPU來說,最大的優勢就是計算並行化,通過向量去計算能夠儘可能地達到並行化的效果。

例如一次矩陣乘法:
c = a1 * b1 + a2 * b2 + a3 * b3 + a4 * b4;
可以改爲
c = dot(vec4(a1, a2, a3, a4), vec4(b1,b2,b3,b4));

向量化的同時也要配合內存佈局的優化;

姿勢二:內存佈局優化

如果讀了上面《利用WebGL2 實現Web前端的GPU計算》這篇文章的同學應該瞭解到,在GPU內所有的數據存儲都是通過Texture的,而Texture本身是一個 長n * 寬m * 通道(rgba)4 的東西,如果我們要存一個3 * 224 * 224 * 150的四維矩陣進去要怎麼辦呢?肯定會涉及到矩陣的編碼,即以一定的格式把高維矩陣存進特性形狀的Texture內,而Texture的數據排布又會影響計算過程中的讀存性能。例如,舉一個較簡單的例子:

如果是常規內存排布的話,計算一次需要按行或者案列遍歷矩陣一次,而GPU的cache是tile類型的,即n*n類型的緩存,根據不同芯片n有所不同。所以這種遍歷方式會頻繁造成cache miss,從而成爲性能的瓶頸。所以,我們就要通過內存排布的方式進行性能優化。類似下圖:

姿勢三:圖優化

由於一個模型是一個一個的算子組成的,而在GPU內每個算子被設計成一個webgl program,每次切換program的時候會造成較多的性能損耗。所以如果有一種手段能夠減少模型的program數量,對性能的提升也是十分可觀的。如下圖:

我們將一些可以融合的節點在圖結構上進行融合(nOP -> 1OP),基於新的計算結點實現新的OP。這樣一來大大減少了OP的數量,進而減少了Program的數量,所以提升了推理性能。在低端手機上效果尤爲明顯。

姿勢四:混合精度計算

以上所有的計算都是基於常規浮點數計算,也就是float32單精度浮點數計算。那麼,在GPU內是否能實現混合精度的計算呢?例如float16、float32、uint8混合精度的計算。答案是可以的,在GPU內實現混合精度計算的價值是在於提升GPU的bandwidth。由於webgl的texture每一個像素點包含rgba四個通道,而每個通道最高爲32位,我們可以在32位內儘可能存儲更多的數據。如果精度爲float16,那麼可以存儲兩個float16,bandwidth就是之前的2倍,同理uint8的bandwidth是之前的4倍。這個性能的提升就是巨大的。還是上圖說話吧:

姿勢n:...

優化的手段還有很多,這裏就不一一列舉了。

引擎落地的場景

目前,基於我們深度優化的引擎已經落地螞蟻集團及阿里經濟體多個應用場景,比較典型的就是文章開頭演示的寵物識別,還有卡證識別、碎屏相機等等等場景。

業界的有之前比較火的虛擬試妝小程序等。

讀到這篇文章的朋友們也可以打開你們的腦洞,挖掘出更多更好玩的智能場景。

未來展望

隨着市面是機型的更新換代及引擎的深入優化,我相信tfjs會在更多富交互的場景上大放異彩,例如擁有AI能力的前端遊戲、AR、VR等等場景。現在我們要做的就是靜下心來,站在巨人的肩膀上持續打磨我們的引擎,願等花開。

作者 | 青壁

原文鏈接

本文爲阿里雲原創內容,未經允許不得轉載。

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