使用 SIMD 和多线程增强 TensorFlow.js WebAssembly 后端

本文最初发表在 TensorFlow 官博，经原作者授权，InfoQ 中文站翻译并分享。

今年三月，我们为 TensorFlow.js 引入了一个新的 WebAssembly（Wasm）加速后端。今天，我们很高兴的宣布一个重大的性能更新：从 TensorFlow.js 2.3.0 版本开始，通过 XNNPACK（一个高度优化的神经网络运算符库）利用 SIMD（向量）指令和多线程，我们的 Wasm 后端速度提高了 10 倍。

基准测试

SIMD 和多线程为我们的 Wasm 后端带来了重大的性能提升。下面是 Google Chrome 浏览器的基准测试，展示了 BlazeFace 的改进。BlazeFace 是一个具有 10 万个参数和大约 2000 万次乘加运算的轻量级模型。

（表中所列出的时间为每次推断的毫秒数）

更大的模型，如 MobileNet V2，这是一个拥有 350 万个参数和大约 3 亿个乘加运算的中性模型，可以获得更快的加速：

注：由于移动浏览器中的多线程支持仍在进行中，故 Pixel 4 不支持 TF.js 多线程 Wasm 后端的基准测试。iOS 版本的 SIMD 支持也仍在开发中。TF.js 多线程 Wasm 后端的节点支持即将推出。

SIMD 和多线程的性能提高是相互独立的。这些基准测试表明，SIMD 给纯 Wasm 带来了 1.7~4.5 倍的性能提升，而多线程在此基础上又带来了 1.8~2.9 倍的加速。

用法

从 TensorFlow.js 2.1.0 开始支持 SIMD，从 TensorFlow.js 2.3.0 开始支持多线程。

在运行时，我们测试 SIMD 和多线程支持，并提供是当的 Wasm 二进制文件。今天，我们为以下每种情况提供不同的二进制代码：

• 默认：运行时不支持 SIIMD 或多线程。
• 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 等 bundler，你可能需要使用我们的 setWasmPaths 帮助其手动指示 Wasm 二进制文件的位置：

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

有关更多信息，请参阅自述文件中的“使用 bundler”一节。

1. 使用脚本标签

<!-- Import @tensorflow/tfjs or @tensorflow/tfjs-core -->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs"></script>


<!-- Adds the WAsm backend to the global backend registry -->
<script src="https://cdn.jsdelivr.net/npm/@tensorflow/tfjs-backend-wasm/dist/tf-backend-wasm.js"></script>


<script>
  tf.setBackend('wasm').then(() => main());
</script

注：TensorFlow.js 为每个后端定义了优先级，并自动为给定的环境选择最佳支持的后端。现在，WebGL 的优先级最高，其次是 Wasm，然后是普通 JS 后端。要想始终使用 Wasm 后端，我们需要显式地调用 tf.setBackedd('wasm')。

演示

要查看性能改进，请查看 BlazeFace 模型的演示，该演示已更新为使用新的 Wasm 后端：https://tfjs-wasm-simd-demo.netlify.app/ 要与未经优化的二进制文件进行比较，请尝试此版本的演示，它手动关闭了 SIMD 和多线程支持。

什么是 Wasm？

WebAssembly（Wasm）是一种跨浏览器的二进制格式，它为 Web 带来了接近原生的代码执行速度。Wasm 可以作为静态类型的高级语言（如 C、C++、Go 和 Rust）编写的程序的编译目标。在 TensorFlow.js 中，我们用 C++实现了 Wasm 后端，并使用 Emscripten 进行编译。XNNPACK 库在下面提供了大量优化的神经网络操作符的实现。

自 2017 年以来，Wash 得到了 Chrome、Safari、Firefox 和 Edge 的支持，全球 90% 的设备都支持 Wasm。

WebAssembly 规范正在快速发展，浏览器正在努力支持越来越多的实验性特性。你可以访问此站点查看你的运行时支持哪些功能，包括：

1. SIMD

SIMD 表示单指令、多数据（Single Instruction, Multiple Data），这意味着 SIMD 指令操作是在固定大小的小元素向量上进行操作的，而不是在单个标量上操作。Wasm SIMD 提案使得现代处理器支持的 SIMD 指令可以在网络浏览器内使用，从而获得显著的性能提升。

Wasm SIMD 是一个第三阶段的建议，并可通过 Chrome 84~86 版本的尝鲜计划（origin trial）获得。这意味着开发人员可以在他们的网站中选择 Wasm SIMD，他们的所有访问者都将能享受到它的好处，而无需在他们的浏览器设置中显式地启用该功能。除了 Gogole Chrome，Firefox Nightly 也默认支持 Wasm SIMD。

2. 多线程

几乎所有的现代处理器都有多个核心，每个核心都能够独立并同时执行指令。WebAssembly 程序可以通过线程建议将它们的工作分散到各个核心以提高性能。这个建议允许不同的 Web 工作器中的多个 Wasm 实例共享单个WebAssembly.Memory 对象，以便在工作器之间进行快速通信。

Wasm 现成是第二阶段的提案，自 74 版本以来，Chrome 桌面版就默认提供了该功能。目前正在进行跨浏览器的工作，以便在移动设备上启用这一功能。

要查看哪些浏览器支持 SIMD、现成和其他实验性特性，请查看 WebAssembly 路线图。

其他改进

自三月份 Wasm 后端最初发布以来，我们已经扩展了运营商的覆盖范围，现在支持 70 多家运营商。许多新的操作符都通过 XNNPACK 库加速，并解锁了对其他模型的支持，如 HandPose 模型。

展望未来

我们期望 Wasm 后端的性能能够不断提高。我们正在密切关注 WebAssembly 中几个不断发展的规范的进展，包括用于更广泛的 SIMD 的灵活向量、准积和熔加运算（quasi fused multiply-add），以及伪最小和最大指令。我们还期待着 ES6 模块对 WebAssembly 模块的支持。与 SIMD 和多线程一样，我们打算在这些功能可用时利用它们，而不影响 TF.js 用户代码。

更多信息

• 查看 WebAssembly 路线图：https://webassembly.org/roadmap/
• 关注 Wasm 规范的进展：https://github.com/WebAssembly/spec
• 欲了解有关 Wasm SIMD 提案的更多信息，请访问：https://github.com/WebAssembly/simd/blob/master/proposals/simd/SIMD.md
• 欲了解有关 Wasm 线程提案的更多信息，请访问：https://github.com/WebAssembly/threads
• 通过 GitHub 上的 issues （问题）和 PR （拉取请求）提交反馈和贡献。
• 使用 #MadeWithTFJS 标签在社交媒体上分享你通过这些优化构建的内容，并加入 TensorFlow.js 社区论坛，以获取未来的更新。

作者介绍：

Ann Yuan 与 Marat Dukhan，均为 Google 的软件工程师。

原文链接：

https://blog.tensorflow.org/2020/09/supercharging-tensorflowjs-webassembly.ht