DEV Community

Cover image for TensorFlow.js 快看不到 LiteRT.js 的車尾燈了
JH5
JH5

Posted on

TensorFlow.js 快看不到 LiteRT.js 的車尾燈了

#tensorflow #machinelearning #javascript #performance

TensorFlow.js 快看不到 LiteRT.js 的車尾燈了

不久前 Chrome更新，也號稱把 Gemini Nano 塞進瀏覽器來讓大家使用， 接著又在月初 Google 官方宣稱 LiteRT.js 配上 WebGPU 後，速度「最高可比 TF.js 快 3 倍」。

說實話，我第一時間是半信半疑的 XD 可能是過去的 PTSD ，經驗上在瀏覽器上跑這些模型都點懷疑人生，不過還是覺得可以來踹踹，結果試玩後讓我有點驚訝。

在 MobileNetV2 上，LiteRT.js 跑 WebGPU backend 的單次推論時間只有 0.41 ms，換算下來是驚人的 2,439 FPS，對比 TensorFlow.js 在 WebGL 下的 10.82 ms（92 FPS），這結果好像根本不止官方講的 3 倍...

接著，我又把模型換成 EfficientNet-Lite4（17M params）時，latency 也只從 0.41ms 稍微上升到 0.62ms，模型增加接近 5 倍，速度居然只慢了 50% ？

LiteRT.js 是什麼

簡單來說，LiteRT.js 就是 LiteRT 的 JavaScript 語言橋樑。

它把原本原生的 C++ runtime 打包編譯成 WebAssembly，並在瀏覽器裡一口氣提供了 WebGPU、WebNN、WASM 三種 backend 選擇。

這不是一個全新的 AI 框架，而是直接承接了龐大的 .tflite 生態系，目前你在 Hugging Face（litert-community）或 Kaggle 上抓下來的大量預訓練模型，通通都能直接拿來用。

測試環境

  • 硬體： M2 Max （30-core GPU, 96GB 統一記憶體）

  • 瀏覽器： Playwright + Chromium（開啟 --enable-webgpu --use-angle=metal，讓 WebGPU 直接走 Metal API）

  • 測試模型：

MobileNetV2 1.0 224（float32，13MB，3.5M params）

EfficientNet-Lite4（float32，49MB，17M params）

  • 流程： 每組先行 Warmup 5 次，接著正式執行 50 次推論，分別記錄 Cold Start（含 Shader 編譯的首次執行）與穩定後的平均 Latency。

MobileNetV2：LiteRT.js WebGPU 0.41ms

第一個結果出來的時候我還懷疑是不是我哪裡搞錯了ＸＤ

LiteRT.js (WebGPU) — MobileNetV2
  Cold start: 4.3 ms
  Average: 0.41 ms
  Min: 0.20 ms / Max: 0.80 ms
  StdDev: 0.12 ms
  Throughput: 2,439 FPS
Enter fullscreen mode Exit fullscreen mode

每次推論平均只需 0.41 毫秒，而且在 50 次測試裡面，最慢的一步也才 0.8ms，標準差只有 0.12ms，表現很穩啊。

有個小細節是，WebGPU 的 model.run() 回傳的是 Promise，resolve 的時間點可能會比 GPU 晶片真正跑完算式稍微早那麼一點點，0.41ms 的數字或許帶有微小的樂觀誤差，但就算我們打到骨折，算它 1ms ，那也是能在 1 秒內處理 1,000 張圖片的驚人效能！

而且 Cold Start 只要 4.3ms，完全不需要像以往 WebGL 那樣經歷痛苦 Cold Start 等待 (PTSD)。

WASM backend：純 CPU 的極限

同一個模型、同一套 runtime，backend 從 webgpu 換成 wasm：

LiteRT.js (WASM / XNNPACK) — MobileNetV2
  Cold start: 18.6 ms
  Average: 13.81 ms
  Min: 13.5 ms / Max: 14.3 ms
  StdDev: 0.17 ms
  Throughput: 72 FPS
Enter fullscreen mode Exit fullscreen mode

13.81ms（約 72 FPS）對絕大多數前端 Vision 應用來說已經非常夠用，而且 WASM 的延遲同樣非常穩定，最大的價值在於絕對的相容性，不用擔心 delpoy 後某個使用者的顯示卡好壞或瀏覽器支援度，不挑硬體。

TensorFlow.js WebGL

再來看看大家最熟悉的舊朋友 TensorFlow.js：

對照組：

TensorFlow.js (WebGL) — MobileNetV2
  Cold start: 10,014 ms（10 秒）
  Average: 10.82 ms
  Min: 10.2 ms / Max: 12.2 ms
  Throughput: 92 FPS
Enter fullscreen mode Exit fullscreen mode

啟動一段時候到穩定後的 10.82ms 其實還算通暢，但那個 Cold Start 等待簡直是災難，WebGL 在第一次執行時必須編譯，整個硬生生卡了十秒才出結果。

雖說 Chrome 會快取編譯好的 Shader，第二次打開頁面就降到 359ms，但只要使用者是第一次造訪或是習慣清 Cache，這 10 秒就是硬傷，相較之下，LiteRT.js WebGPU 無論第幾次開啟，Cold Start 都穩穩維持在 4.3ms。

三個 backend 放在一起比

Framework Avg Latency FPS Cold Start Model Load
LiteRT.js (WebGPU) 0.41 ms 2,439 FPS 4.3 ms 143 ms
LiteRT.js (WASM) 13.81 ms 72 FPS 18.6 ms 52 ms
TensorFlow.js (WebGL) 10.82 ms 92 FPS 10,014 ms 1,516 ms

為什麼 WebGPU 能把 WebGL 壓在地上打？因為兩者的基因根本不同。WebGL 是硬把神經網路運算塞進 3D 圖形的 Fragment Shader 裡，而 WebGPU 從第一天起就是為了通用計算（Compute Shader）與 ML 推論設計的。

模型測試：EfficientNet-Lite4

小模型表現好不稀奇，那換成大模型呢 （跟現有LLM參數量比是超迷你模型...）？

我抓了 EfficientNet-Lite4（17M 參數量、49MB 檔大小、300x300 輸入，Top-1 Accuracy 80.4%），參數量大概是 MobileNetV2 的4倍多：

EfficientNet-Lite4 (WebGPU)
  Cold start: 0.8 ms
  Average: 0.62 ms
  Min: 0.40 ms / Max: 1.10 ms
  StdDev: 0.15 ms
  Throughput: 1,623 FPS
Enter fullscreen mode Exit fullscreen mode

結果很讓人吃驚：0.62ms，依舊能維持 1,623 FPS

兩顆模型 WebGPU 並排：

Metric MobileNetV2 EfficientNet-Lite4 Delta
Params 3.5M 17M 4.85x
WebGPU latency 0.41 ms 0.62 ms 1.51x
FPS 2,439 1,623 0.67x
Cold start 4.3 ms 0.8 ms 更快

參數量多出近 5 倍，運算延遲竟然只微幅增加 50%，這是因為在 GPU 上的計算密集度隨 Channel 增加而提高，反而更能吃滿 GPU 的平行運算核心的優勢啊。

WASM 在大模型上ＧＧ

把相同的 EfficientNet-Lite4 丟給 WASM backend，結果直接就悲劇了

EfficientNet-Lite4 (WASM / XNNPACK)
  Cold start: 131.8 ms
  Average: 124.37 ms
  Min: 121.90 ms / Max: 127.40 ms
  StdDev: 1.36 ms
  Throughput: 8 FPS
Enter fullscreen mode Exit fullscreen mode

124ms，每秒 8 張。WASM 直接崩了。

Model WASM Latency FPS
MobileNetV2 13.81 ms 72
EfficientNet-Lite4 124.37 ms 8
Scale factor 9.01x

參數量多 4.85 倍，WASM latency 多了 9 倍——latency 幾乎跟模型大小線性成長，沒有 GPU 那種平行化紅利。

加速比的對比更驚人：

Model WebGPU WASM Speedup
MobileNetV2 0.41 ms 13.81 ms 33.7x
EfficientNet-Lite4 0.62 ms 124.37 ms 200.6x

如果你的應用場景現在還是只能base WASM，你的模型選擇會被死死限制在 MobileNetV2 這種輕量級規模，不過一啖可直上了 WebGPU，你幾乎可以放飛自我～

聊聊實務限制

優點講完了，接著聊聊我在實測過程中踩到的坑與一些現實限制：

  1. Batch inference 的限制

我本來想測試一次餵 2 或 4 張圖片進行 Batch 推論，結果程式直接噴錯，應該是因爲在於現成下載的 .tflite 模型在 Export 時就把 Batch Dimension 鎖死在 [1, 224, 224, 3] 了，如果想要在前端做批次處理，不能像 Python 那樣隨意傳入 Tensor，必須自己透過 Multiple Model Instances 或 Web Worker Pool 來平行處理。

  1. 手動記憶體管理

LiteRT.js 必須手動呼叫 .delete() 來釋放 Tensor，前面剛開始試的時候踩了幾次洞，後來發現只要有落實 cleanup，記憶體會穩定維持在 16~18MB 之間，完全沒有 Leak，這部分應該之後會有貼心的patch吧。

  1. 行動端與瀏覽器相容性 目前 WebGPU 在 Mobile 瀏覽器上的支援度與穩定度依然是尚未開發國家階段，Google 官方目前也不太推薦在手機端硬開 WebGPU，此外，Safari 目前對 WebGPU 仍預設為實驗性功能，要等到普及應該還需要不少時間。

該從 TF.js 跳過來嗎

  1. 如果你的應用目前主力使用者是 Desktop Chrome / Firefox
    那就別猶豫了，現在就上吧！

  2. 如果產品需要兼顧 Safari 或手機網頁：
    可以先架構遷移至 LiteRT.js，但預設先走 WASM Backend（ MobileNetV2 跑 13.8ms 依舊很順暢），等未來 Safari 正式全面解鎖 WebGPU 時，就能無縫升級爆發力。

  3. 從生態系發展來看：
    Google 目前的重心顯然已經轉向 LiteRT，TensorFlow.js 的更新維護速度已經明顯放緩，與其等舊架構被淘汰，不如趁現在開始規劃。

Top comments (0)