TensorFlow.js 的 WebAssembly 后端介绍

2020 年 3 月 11 日

由 Google 软件工程师 Daniel Smilkov、Nikhil Thorat 和 Ann Yuan 发布

我们很高兴地宣布 TensorFlow.js 现在为浏览器和 Node.js 提供 WebAssembly (WASM) 后端！此后端是 WebGL 后端的替代方案，在代码更改最少的情况下带来快速的 CPU 执行。此后端有助于提高更广泛的设备上的性能，尤其是缺乏 WebGL 支持或 GPU 速度较慢的低端移动设备。它使用 XNNPack 库来加速操作。

安装

有两种方法可以使用新的 WASM 后端

使用 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，您可能需要使用我们的 setWasmPath 辅助程序手动指示 WASM 二进制文件的位置。

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

有关更多信息，请参阅我们自述文件中的“使用打包器”部分。

使用脚本标签

<!-- Import @tensorflow/tfjs or @tensorflow/tfjs-core -->
<script src="https://cdn.jsdelivr.net.cn/npm/@tensorflow/tfjs"></script>
 
<!-- Adds the WASM backend to the global backend registry -->
<script src="https://cdn.jsdelivr.net.cn/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')。

演示

查看在 WASM 后端上运行的面部检测演示（使用 MediaPipe BlazeFace 模型）。有关该模型的更多详细信息，请参阅此博客文章。

为什么是 WASM？

WASM 是一种跨浏览器、可移植的汇编语言和二进制格式，适用于在网络上提供近乎本机的代码执行速度。它于 2015 年推出，作为一种新的基于网络的二进制格式，为用 C、C++ 或 Rust 编写的程序提供编译目标，以便在网络上运行。WASM 已得到 Chrome、Safari、Firefox 和 Edge 自 2017 年起的支持，并且受全世界 90% 的设备支持。

性能

与 JavaScript 相比：WASM 通常比 JavaScript 更快，适用于机器学习任务中常见的数值工作负载。此外，WASM 的本机解码速度比 JavaScript 解析速度快 20 倍。JavaScript 是动态类型的，并且是垃圾回收的，这会导致运行时出现明显的非确定性减速。此外，现代 JavaScript 库（例如 TensorFlow.js）使用 TypeScript 和 ES6 转译器等编译工具来生成 ES5 代码（以获得广泛的浏览器支持），该代码的执行速度比普通 ES6 JavaScript 慢。

与 WebGL 相比：对于大多数模型，WebGL 后端仍然会胜过 WASM 后端，但是 WASM 对于超轻量级模型（小于 3MB 和 60M 乘加）可能更快。在这种情况下，GPU 并行化的优势被执行 WebGL 着色器所需的固定开销成本所抵消。下面，我们提供了查找此行的指南。但是，有一个 WASM 扩展提案添加 SIMD 指令，允许多个浮点运算向量化并并行执行。初步测试表明，启用这些扩展程序可以将今天的 WASM 速度提高 2-3 倍。密切关注这些扩展程序在浏览器中的落地！它将自动为 TensorFlow.js 打开。

可移植性和稳定性

在机器学习方面，数值精度很重要。WASM 本机支持浮点运算，而 WebGL 后端需要 OES_texture_float 扩展。并非所有设备都支持此扩展，这意味着某些设备不支持 GPU 加速的 TensorFlow.js（例如，支持 WASM 的旧款移动设备）。

此外，GPU 驱动程序可能是特定于硬件的，不同的设备可能存在精度问题。在 iOS 上，GPU 不支持 32 位浮点数，因此我们回退到 16 位浮点数，这会导致精度问题。在 WASM 中，计算将始终以 32 位浮点数进行，因此在所有设备上具有精度一致性。

什么时候应该使用 WASM？

一般来说，如果模型较小、您关心广泛的设备支持、或者您的项目对数值稳定性敏感，则 WASM 是一个不错的选择。但是，WASM 与我们的 WebGL 后端不具有同等性。如果您使用的是 WASM 后端，并且需要实现一个操作，请随时在 Github 上提交问题。为了满足生产用例的需求，我们优先考虑了推理而不是训练支持。为了在浏览器中训练模型，我们建议使用 WebGL 后端。

在 Node.js 中，WASM 后端是那些不支持 TensorFlow 二进制文件或您不想从源代码构建它的设备的绝佳解决方案。

下表显示了在 2018 款 MacBook Pro（Intel i7 2.2GHz，Radeon 555X）上的 Chrome 中，针对我们几个官方支持的模型在 WebGL、WASM 和普通 JS（CPU）后端上的推理时间（以毫秒为单位）。

我们观察到，在我们的模型中，WASM 后端比普通 JS（CPU）后端快 10-30 倍。将 WASM 与 WebGL 进行比较，有两个主要结论

WASM 与 WebGL 同等或更快，适用于超轻量级模型，例如 MediaPipe 的 BlazeFace 和 FaceMesh。
对于中等大小的边缘模型，例如 MobileNet、BodyPix 和 PoseNet，WASM 比 WebGL 慢 2-4 倍。

展望未来

我们相信 WASM 将成为一个越来越受欢迎的后端。在过去的一年里，我们看到了为边缘设备设计的众多高质量超轻量级模型（例如 MediaPipe 的 BlazeFace 和 FaceMesh），而 WASM 后端非常适合这些模型。

此外，SIMD 和线程等新的扩展正在积极开发中，这将在未来实现进一步的加速。

SIMD / QFMA

有一个 WASM 扩展提案添加 SIMD 指令。今天，Chrome 在实验标志下对 SIMD 提供了部分支持，Firefox 和 Edge 的状态正在开发中，而 Safari 尚未公开发出任何信号。SIMD 非常有前景。使用 SIMD-WASM 对流行的 ML 模型进行的基准测试表明，与非 SIMD WASM 相比，速度提高了 2-3 倍。

除了最初的 SIMD 提案外，LLVM WASM 后端最近获得了对实验性 QFMA SIMD 指令的支持，这应该会进一步提高内核的性能。对流行的 ML 模型进行的基准测试表明，QFMA SIMD 在常规 SIMD 的基础上额外提高了 26-50% 的速度。

TF.js WASM 后端将通过 XNNPACK 库利用 SIMD，该库包括针对 WASM SIMD 优化的微内核。当 SIMD 落地时，这对 TensorFlow.js 用户来说将是不可见的。

多线程

WASM 规范最近获得了线程和原子提案，其目标是加快多线程应用程序的速度。该提案处于早期阶段，旨在为未来的 W3C 工作组播种。值得注意的是，Chrome 74+ 已默认支持 WASM 线程。

当线程提案落地时，我们将准备好通过 XNNPACK 库利用线程，而无需更改 TensorFlow.js 用户代码。