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

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

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

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

基準(zhǔn)

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

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

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

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

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

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

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

SIMD 和多線程帶來(lái)的性能提升彼此獨(dú)立。這些基準(zhǔn)表明，SIMD 將標(biāo)準(zhǔn) Wasm 的性能提高了 1.7-4.5 倍，而多線程在此基礎(chǔ)上又帶來(lái)了 1.8-2.9 倍的速度提升。

用法

SIMD 從 TensorFlow.js 2.1.0 開(kāi)始得到支持，多線程從 TensorFlow.js 2.3.0 開(kāi)始得到支持。

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

默認(rèn)：運(yùn)行時(shí)不支持 SIMD 或多線程

SIMD：運(yùn)行時(shí)支持 SIMD，但不支持多線程

SIMD + 多線程：運(yùn)行時(shí)支持 SIMD 和多線程

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

1. 通過(guò) 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());

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

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

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

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

2. 通過(guò)腳本標(biāo)記

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

演示

要實(shí)際感受性能提升，可以來(lái)看看我們的 BlazeFace 模型演示。該模型已更新為使用新的 Wasm 后端：https://tfjs-wasm-simd-demo.netlify.app/。要與未優(yōu)化的二進(jìn)制文件進(jìn)行比較，請(qǐng)?jiān)囋嚧税姹狙菔荆梢允謩?dòng)關(guān)閉 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) 是一種跨瀏覽器的二進(jìn)制文件格式，為網(wǎng)絡(luò)帶來(lái)了接近原生的代碼執(zhí)行速度。Wasm 可以作為 C、C++、Go 和 Rust 等靜態(tài)類型高級(jí)語(yǔ)言所編寫(xiě)程序的編譯目標(biāo)。在 TensorFlow.js 中，我們使用 C++ 實(shí)現(xiàn) Wasm 后端，并使用 Emscripten 編譯。XNNPACK 庫(kù)在下方提供了神經(jīng)網(wǎng)絡(luò)算子的高度優(yōu)化實(shí)現(xiàn)。

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

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

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

Wasm SIMD 是一個(gè)第 3 期提案，通過(guò) Chrome 84-86 中的初始試用提供。這意味著開(kāi)發(fā)者可以在網(wǎng)站上選擇使用 Wasm SIMD，讓所有訪問(wèn)者直接獲益，而無(wú)需在瀏覽器設(shè)置中明確啟用該功能。除了 Google Chrome，F(xiàn)irefox Nightly 也默認(rèn)支持 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)代處理器都有多個(gè)核心，每個(gè)核心都能獨(dú)立并發(fā)地執(zhí)行指令。WebAssembly 程序可以通過(guò)線程提案將工作分布到不同核心，進(jìn)而提高性能。這個(gè)提案允許多個(gè) Wasm 實(shí)例在不同的 Web 工作者中共享一個(gè) WebAssembly.Memory 對(duì)象，實(shí)現(xiàn)工作進(jì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 版起默認(rèn)可用。為了在移動(dòng)設(shè)備實(shí)現(xiàn)這一功能，跨瀏覽器的開(kāi)發(fā)工作也已開(kāi)始。

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

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

WebAssembly 路線圖顯示了支持 SIMD、線程和其他實(shí)驗(yàn)功能的瀏覽器。

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

其他改進(jìn)

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

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

展望未來(lái)

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

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

準(zhǔn)融合乘加
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/

關(guān)注 Wasm 規(guī)范的進(jìn)展

https://github.com/WebAssembly/spec

詳細(xì)了解 Wasm SIMD 提案

https://github.com/WebAssembly/simd

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

通過(guò) GitHub 上提 issue 和 PR 提交反饋和貢獻(xiàn)
https://github.com/tensorflow/tfjs/issues/new
https://github.com/tensorflow/tfjs/pulls

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

致謝

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

責(zé)任編輯：xj

原文標(biāo)題：SIMD 和多線程大幅增強(qiáng) TFJS WebAssembly 后端

文章出處：【微信公眾號(hào)：TensorFlow】歡迎添加關(guān)注！文章轉(zhuǎn)載請(qǐng)注明出處。