重大性能更新：Wasm 后端將利用 SIMD指令和 XNNPACK多線程

3 月，我們?yōu)?TensorFlow.js 推出了一個新的 WebAssembly(Wasm) 加速后端（繼續(xù)閱讀以進一步了解 Wasm 及其重要性）。今天，我們很高興宣布一項重大性能更新：自 TensorFlow.js 版本 2.3.0 起，我們的 Wasm 后端將利用 SIMD（向量）指令和 XNNPACK（一種高度優(yōu)化的神經網(wǎng)絡算子庫）多線程實現(xiàn) 10 倍提速。

SIMD（向量）指令
https://github.com/WebAssembly/simd

XNNPACK
https://github.com/google/XNNPACK

多線程
https://github.com/WebAssembly/threads

基準

SIMD 和多線程為我們的 Wasm 后端帶來重大性能提升。BlazeFace是擁有 10 萬個參數(shù)和大約 2000 萬次乘加運算的輕型模型。以下是在 Google Chrome 瀏覽器中演示了 BlazeFace 的性能評測：

（所列時間為每次推理的毫秒數(shù)）

對于更大的模型，如擁有 350 萬個參數(shù)和大約 3 億次乘加運算的中型模型 MobileNet V2，加速效果會更加明顯：

MobileNet V2
https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/classification/2

*注：由于移動瀏覽器中的多線程支持仍在開發(fā)中，因此 Pixel 4 無法使用 TF.js 多線程 Wasm 后端基準。iOS 中的 SIMD 支持也仍處于開發(fā)階段。

移動瀏覽器中的多線程支持仍在開發(fā)中
https://www.chromestatus.com/feature/5724132452859904

**注：我們即將推出 TF.js 多線程 Wasm 后端的節(jié)點支持。

SIMD 和多線程帶來的性能提升彼此獨立。這些基準表明，SIMD 將標準 Wasm 的性能提高了 1.7-4.5 倍，而多線程在此基礎上又帶來了 1.8-2.9 倍的速度提升。

用法

SIMD 從 TensorFlow.js 2.1.0 開始得到支持，多線程從 TensorFlow.js 2.3.0 開始得到支持。

在 運行時 (Runtime)，我們測試 SIMD 和多線程支持并提供適當?shù)?Wasm 二進制文件。今天，我們?yōu)橐韵旅糠N情況提供不同的二進制文件：

默認：運行時不支持 SIMD 或多線程

SIMD：運行時支持 SIMD，但不支持多線程

SIMD + 多線程：運行時支持 SIMD 和多線程

由于大多數(shù)支持多線程的運行時也會支持 SIMD，因此我們決定忽略僅支持多線程的運行時，縮減軟件包的大小。如果您的運行時支持多線程而不支持 SIMD，您將獲得默認的二進制文件。您可以通過兩種方式使用 Wasm 后端：

1. 通過 NPM

// Import @tensorflow/tfjs or @tensorflow/tfjs-core const tf = require('@tensorflow/tfjs'); // Add the WAsm backend to the global backend registry. require('@tensorflow/tfjs-backend-wasm'); // Set the backend to WAsm and wait for the module to be ready. tf.setBackend('wasm').then(() => main());

此庫預期 Wasm 二進制文件相對于主 JS 文件定位。如果您使用的是 Parcel 或 Webpack 等打包工具，則可能需要使用我們的 setWasmPaths 幫助程序手動指示 Wasm 二進制文件的位置：

import {setWasmPaths} from '@tensorflow/tfjs-backend-wasm'; setWasmPaths(yourCustomFolder);tf.setBackend('wasm').then(() => {...});

請參閱我們 README 上的“使用打包工具”部分了解詳情。

使用打包工具
https://github.com/tensorflow/tfjs/tree/master/tfjs-backend-wasm#using-bundlers

2. 通過腳本標記

注：TensorFlow.js 為每個后端定義一個優(yōu)先級，并將針對給定環(huán)境條件自動選擇最受支持的后端?，F(xiàn)在，WebGL 具有最高優(yōu)先級，其次是 Wasm，然后是普通 JS 后端。要始終使用 Wasm 后端，我們需要顯式調用 tf.setBackend(‘wasm’)。

演示

要實際感受性能提升，可以來看看我們的 BlazeFace 模型演示。該模型已更新為使用新的 Wasm 后端：https://tfjs-wasm-simd-demo.netlify.app/。要與未優(yōu)化的二進制文件進行比較，請試試此版本演示，可以手動關閉 SIMD 和多線程支持。

此版本演示
https://storage.googleapis.com/tfjs-models/demos/blazeface/index.html?tfjsflags=WASM_HAS_MULTITHREAD_SUPPORT:false,WASM_HAS_SIMD_SUPPORT:false

什么是 Wasm？

WebAssembly (Wasm) 是一種跨瀏覽器的二進制文件格式，為網(wǎng)絡帶來了接近原生的代碼執(zhí)行速度。Wasm 可以作為 C、C++、Go 和 Rust 等靜態(tài)類型高級語言所編寫程序的編譯目標。在 TensorFlow.js 中，我們使用 C++ 實現(xiàn) Wasm 后端，并使用 Emscripten 編譯。XNNPACK 庫在下方提供了神經網(wǎng)絡算子的高度優(yōu)化實現(xiàn)。

自 2017 年以來，Wasm 已獲得 Chrome、Safari、Firefox 和 Edge 的支持，并已得到全球 90% 設備的支持。

WebAssembly 規(guī)范發(fā)展迅速，瀏覽器正在盡全力支持越來越多的實驗性功能。您可以訪問此網(wǎng)站查看您的運行時支持哪些功能，包括：

1. SIMD
SIMD 代表 Single Instruction, Multiple Data，這意味著 SIMD 指令是在固定大小的小元素向量而不是各個標量上進行運算。Wasm SIMD 提案使現(xiàn)代處理器支持的 SIMD 指令可以在網(wǎng)絡瀏覽器內使用，達成顯著的性能提升。

Wasm SIMD 是一個第 3 期提案，通過 Chrome 84-86 中的初始試用提供。這意味著開發(fā)者可以在網(wǎng)站上選擇使用 Wasm SIMD，讓所有訪問者直接獲益，而無需在瀏覽器設置中明確啟用該功能。除了 Google Chrome，F(xiàn)irefox Nightly 也默認支持 Wasm SIMD。

第 3 期
https://github.com/WebAssembly/meetings/blob/master/process/phases.md#3-implementation-phase-community--working-group

初始試用
https://developers.chrome.com/origintrials/#/view_trial/-4708513410415853567

2. 多線程
幾乎所有現(xiàn)代處理器都有多個核心，每個核心都能獨立并發(fā)地執(zhí)行指令。WebAssembly 程序可以通過線程提案將工作分布到不同核心，進而提高性能。這個提案允許多個 Wasm 實例在不同的 Web 工作者中共享一個 WebAssembly.Memory 對象，實現(xiàn)工作進程之間的快速通信。

線程提案
https://github.com/WebAssembly/threads/blob/master/proposals/threads/Overview.md

WebAssembly.Memory
https://developer.mozilla.org/en-US/docs/Web/JavaScript/Reference/Global_Objects/WebAssembly/Memory

Wasm 線程是第 2 期提案，已在桌面版 Chrome 中自 74 版起默認可用。為了在移動設備實現(xiàn)這一功能，跨瀏覽器的開發(fā)工作也已開始。

第 2 期
https://github.com/WebAssembly/proposals/issues/14

桌面版 Chrome
https://www.chromestatus.com/feature/5724132452859904

WebAssembly 路線圖顯示了支持 SIMD、線程和其他實驗功能的瀏覽器。

路線圖
https://webassembly.org/roadmap/

其他改進

自從 3 月 Wasm 后端的首次發(fā)布，我們已經擴大了算子的覆蓋范圍，現(xiàn)在支持超過 70 種算子。許多新的算子都是通過 XNNPACK 庫加速，并解鎖了對額外模型的支持，比如 HandPose 模型。

HandPose
https://github.com/tensorflow/tfjs-models/tree/master/handpose

展望未來

我們希望不斷提高 Wasm 后端的性能。我們正在密切關注 WebAssembly 中不斷發(fā)展的規(guī)范，包括用于更廣泛 SIMD 的靈活向量，準融合乘加，以及偽最小和最大指令。我們也期待著 ES6 模塊對 WebAssembly 模塊的支持。與 SIMD 和多線程一樣，我們打算在這些功能可用時充分加以利用，而不影響 TF.js 用戶代碼。

靈活向量
https://github.com/WebAssembly/flexible-vectors

準融合乘加
https://github.com/WebAssembly/simd/pull/79

偽最小和最大指令
https://github.com/WebAssembly/simd/pull/122

ES6 模塊
https://github.com/WebAssembly/esm-integration

更多信息

查看 WebAssembly 路線圖
https://webassembly.org/roadmap/

關注 Wasm 規(guī)范的進展

https://github.com/WebAssembly/spec

詳細了解 Wasm SIMD 提案

https://github.com/WebAssembly/simd

詳細了解 Wasm 線程提案
https://github.com/WebAssembly/threads

通過 GitHub 上提 issue 和 PR 提交反饋和貢獻
https://github.com/tensorflow/tfjs/issues/new
https://github.com/tensorflow/tfjs/pulls

加入 TensorFlow.js 社區(qū)論壇了解產品更新信息
https://groups.google.com/a/tensorflow.org/g/tfjs

致謝

我們要感謝 Daniel Smilkov 和 Nikhil Thorat 為 WebAssembly 后端和 XNNPACK 集成奠定基礎，感謝 Matsvei Zhdanovich 收集 Pixel 4 基準數(shù)據(jù)，感謝 Frank Barchard 在 XNNPACK 中實現(xiàn)低級 Wasm SIMD 優(yōu)化。

責任編輯：xj

原文標題：SIMD 和多線程大幅增強 TFJS WebAssembly 后端

文章出處：【微信公眾號：TensorFlow】歡迎添加關注！文章轉載請注明出處。