3 月,我们为 TensorFlow.js 推出了一个新的 WebAssembly (Wasm) 加速后端(继续阅读以进一步了解 Wasm 及其重要性)。今天,我们很高兴宣布一项重大性能更新:自 TensorFlow.js 版本 2.3.0 起,我们的 Wasm 后端将利用 SIMD(向量)指令和 XNNPACK(一种高度优化的神经网络算子库)多线程实现 10 倍提速。
SIMD(向量)指令
https://github.com/WebAssembly/simd
XNNPACK
https://github.com/google/XNNPACK
多线程
https://github.com/WebAssembly/threads
基准
SIMD 和多线程为我们的 Wasm 后端带来重大性能提升。BlazeFace 是拥有 10 万个参数和大约 2000 万次乘加运算的轻型模型。以下是在 Google Chrome 浏览器中演示了 BlazeFace 的性能评测:
(所列时间为每次推理的毫秒数)
对于更大的模型,如拥有 350 万个参数和大约 3 亿次乘加运算的中型模型 MobileNet V2,加速效果会更加明显:
MobileNet V2
https://tfhub.dev/google/imagenet/mobilenet_v2_100_224/classification/2
*注:由于移动浏览器中的多线程支持仍在开发中,因此 Pixel 4 无法使用 TF.js 多线程 Wasm 后端基准。iOS 中的 SIMD 支持也仍处于开发阶段。
移动浏览器中的多线程支持仍在开发中
https://www.chromestatus.com/feature/5724132452859904
**注:我们即将推出 TF.js 多线程 Wasm 后端的节点支持。
SIMD 和多线程带来的性能提升彼此独立。这些基准表明,SIMD 将标准 Wasm 的性能提高了 1.7-4.5 倍,而多线程在此基础上又带来了 1.8-2.9 倍的速度提升。
用法
SIMD 从 TensorFlow.js 2.1.0 开始得到支持,多线程从 TensorFlow.js 2.3.0 开始得到支持。
在 运行时 (Runtime),我们测试 SIMD 和多线程支持并提供适当的 Wasm 二进制文件。今天,我们为以下每种情况提供不同的二进制文件:
默认:运行时不支持 SIMD 或多线程
SIMD:运行时支持 SIMD,但不支持多线程
SIMD + 多线程:运行时支持 SIMD 和多线程
由于大多数支持多线程的运行时也会支持 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 为每个后端定义一个优先级,并将针对给定环境条件自动选择最受支持的后端。现在,WebGL 具有最高优先级,其次是 Wasm,然后是普通 JS 后端。要始终使用 Wasm 后端,我们需要显式调用 tf.setBackend(‘wasm’)。
演示
要实际感受性能提升,可以来看看我们的 BlazeFace 模型演示。该模型已更新为使用新的 Wasm 后端:https://tfjs-wasm-simd-demo.netlify.app/。要与未优化的二进制文件进行比较,请试试此版本演示,可以手动关闭 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) 是一种跨浏览器的二进制文件格式,为网络带来了接近原生的代码执行速度。Wasm 可以作为 C、C++、Go 和 Rust 等静态类型高级语言所编写程序的编译目标。在 TensorFlow.js 中,我们使用 C++ 实现 Wasm 后端,并使用 Emscripten 编译。XNNPACK 库在下方提供了神经网络算子的高度优化实现。
自 2017 年以来,Wasm 已获得 Chrome、Safari、Firefox 和 Edge 的支持,并已得到全球 90% 设备的支持。
WebAssembly 规范发展迅速,浏览器正在尽全力支持越来越多的实验性功能。您可以访问此网站查看您的运行时支持哪些功能,包括:
1. SIMD
SIMD 代表 Single Instruction, Multiple Data,这意味着 SIMD 指令是在固定大小的小元素向量而不是各个标量上进行运算。Wasm SIMD 提案使现代处理器支持的 SIMD 指令可以在网络浏览器内使用,达成显著的性能提升。
Wasm SIMD 是一个第 3 期提案,通过 Chrome 84-86 中的初始试用提供。这意味着开发者可以在网站上选择使用 Wasm SIMD,让所有访问者直接获益,而无需在浏览器设置中明确启用该功能。除了 Google Chrome,Firefox 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. 多线程
几乎所有现代处理器都有多个核心,每个核心都能独立并发地执行指令。WebAssembly 程序可以通过线程提案将工作分布到不同核心,进而提高性能。这个提案允许多个 Wasm 实例在不同的 Web 工作者中共享一个 WebAssembly.Memory 对象,实现工作进程之间的快速通信。
线程提案
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 版起默认可用。为了在移动设备实现这一功能,跨浏览器的开发工作也已开始。
第 2 期
https://github.com/WebAssembly/proposals/issues/14
桌面版 Chrome
https://www.chromestatus.com/feature/5724132452859904
WebAssembly 路线图显示了支持 SIMD、线程和其他实验功能的浏览器。
路线图
https://webassembly.org/roadmap/
其他改进
自从 3 月 Wasm 后端的首次发布,我们已经扩大了算子的覆盖范围,现在支持超过 70 种算子。许多新的算子都是通过 XNNPACK 库加速,并解锁了对额外模型的支持,比如 HandPose 模型。
HandPose
https://github.com/tensorflow/tfjs-models/tree/master/handpose
展望未来
我们希望不断提高 Wasm 后端的性能。我们正在密切关注 WebAssembly 中不断发展的规范,包括用于更广泛 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 规范的进展
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 社区论坛了解产品更新信息
https://groups.google.com/a/tensorflow.org/g/tfjs
致谢
我们要感谢 Daniel Smilkov 和 Nikhil Thorat 为 WebAssembly 后端和 XNNPACK 集成奠定基础,感谢 Matsvei Zhdanovich 收集 Pixel 4 基准数据,感谢 Frank Barchard 在 XNNPACK 中实现低级 Wasm SIMD 优化。
责任编辑:xj
原文标题:SIMD 和多线程大幅增强 TFJS WebAssembly 后端
文章出处:【微信公众号:TensorFlow】欢迎添加关注!文章转载请注明出处。
全部0条评论
快来发表一下你的评论吧 !