SIMD 和多线程大幅增强 TFJS WebAssembly 后端

文 / 软件工程师 Ann Yuan 和 Marat Dukhan，Google

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. 通过脚本标记

<!-- 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.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 优化。

—推荐阅读—

了解更多请点击 “阅读原文” 访问官网。