重大性能更新:Wasm 后端将利用 SIMD指令和 XNNPACK多线程

描述

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 的性能评测:

(所列时间为每次推理的毫秒数)

SIMD

对于更大的模型,如拥有 350 万个参数和大约 3 亿次乘加运算的中型模型 MobileNet V2,加速效果会更加明显:

SIMD

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

SIMD

用法

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】欢迎添加关注!文章转载请注明出处。

打开APP阅读更多精彩内容
声明:本文内容及配图由入驻作者撰写或者入驻合作网站授权转载。文章观点仅代表作者本人,不代表电子发烧友网立场。文章及其配图仅供工程师学习之用,如有内容侵权或者其他违规问题,请联系本站处理。 举报投诉

全部0条评论

快来发表一下你的评论吧 !

×
20
完善资料,
赚取积分