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

技术分享
,
1

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

3 月，我们为 TensorFlow.js 推出了一个新的 WebAssembly (Wasm) 加速后端（继续阅读以进一步了解 Wasm 及其重要性）。今天，我们很高兴宣布一项重大性能更新：自 TensorFlow.js 版本 2.3.0 起，我们的 Wasm 后端将利用 SIMD（向量）指令 和 XNNPACK（一种高度优化的神经网络算子库）多线程 实现 10 倍提速。

性能基准测试

SIMD 和多线程为我们的 Wasm 后端带来重大性能提升。BlazeFace 是拥有 10 万个参数和大约 2000 万次乘加运算的轻型模型。以下是在 Google Chrome 浏览器中演示了 BlazeFace 的性能评测：

（所列时间为每次推理的毫秒数）

image
image1380×552 24.5 KB

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

image
image1380×550 25.5 KB

*注：由于 移动浏览器中的多线程支持仍在开发中，因此 Pixel 4 无法使用 TF.js 多线程 Wasm 后端基准。iOS 中的 SIMD 支持也仍处于开发阶段。

**注：我们即将推出 TF.js 多线程 Wasm 后端的节点支持。

SIMD 和多线程带来的性能提升彼此独立。这些基准表明，SIMD 将标准 Wasm 的性能提高了 1.7-4.5 倍，而多线程在此基础上又带来了 1.8-2.9 倍的速度提升。

image
image1200×742 21.8 KB

image
image1200×742 20.1 KB

用法

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 上的“使用打包工具”部分了解详情。

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 和多线程支持。

什么是 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。

2. 多线程
几乎所有现代处理器都有多个核心，每个核心都能独立并发地执行指令。WebAssembly 程序可以通过 线程提案 将工作分布到不同核心，进而提高性能。这个提案允许多个 Wasm 实例在不同的 Web 工作者中共享一个 WebAssembly.Memory 对象，实现工作进程之间的快速通信。

Wasm 线程是 第 2 期提案，已在 桌面版 Chrome 中自 74 版起默认可用。为了在移动设备实现这一功能，跨浏览器的开发工作也已开始。

WebAssembly 路线图 显示了支持 SIMD、线程和其他实验功能的浏览器。

其他改进

自从 3 月 Wasm 后端的首次发布，我们已经扩大了算子的覆盖范围，现在支持超过 70 种算子。许多新的算子都是通过 XNNPACK 库加速，并解锁了对额外模型的支持，比如 HandPose 模型。

展望未来

我们希望不断提高 Wasm 后端的性能。我们正在密切关注 WebAssembly 中不断发展的规范，包括用于更广泛 SIMD 的 灵活向量，准融合乘加，以及 伪最小和最大指令。我们也期待着 ES6 模块 对 WebAssembly 模块的支持。与 SIMD 和多线程一样，我们打算在这些功能可用时充分加以利用，而不影响 TF.js 用户代码。

更多信息

致谢

我们要感谢 Daniel Smilkov 和 Nikhil Thorat 为 WebAssembly 后端和 XNNPACK 集成奠定基础，感谢 Matsvei Zhdanovich 收集 Pixel 4 基准数据，感谢 Frank Barchard 在 XNNPACK 中实现低级 Wasm SIMD 优化。

image

原文：Supercharging the TensorFlow.js WebAssembly backend with SIMD and multi-threading
中文：TensorFlow 公众号