WebAssembly 在 TensorFlow.js 中的应用

译者注：这是《WebAssembly 权威指南》一书的第九章，这一章是关于 TensorFlow.js 的 WebAssembly (WASM) 后端的介绍和使用教程。TensorFlow.js 是一个用于机器学习的 JavaScript 框架，它可以在浏览器和 Node.js 中运行。WASM 后端是一个可选的后端，它可以提高 CPU 的执行速度，而且只需要少量的代码修改。WASM 后端可以帮助提高在更广泛的设备上的性能，特别是那些缺乏 WebGL 支持或 GPU 较慢的低端移动设备。

这是“WebAssembly 应用”的第一章，我在其中重点介绍了该技术的潜在使用场景内。正如你在本书中看到的那样，没有单一的目的。相反，设计师们精心打造了一个平台，其范围已经扩大到几乎触及软件行业的各个方面。所以跟随我的脚步。我不会自己教授新功能。相反，我希望帮助你了解我们行业的变化有多快，以及 WebAssembly 可以提供哪些帮助。

首先，我希望你停下来想想编程语言和机器学习。你想到的第一个编程语言是什么？你的答案很可能不是 JavaScript。为什么会这样？机器学习是一项以性能为导向的活动，因为它具有巨大的计算工作量。

像 Python 这样的语言与机器学习的联系比 JavaScript 强得多。老实说，这也有点牵强。Python 本身就是一种可怕的语言。但是，它在可读性、灵活的编程风格（函数式、面向对象、过程式）以及广泛的算法、可视化和数据操作功能之间取得了良好的平衡。如果你能让它运行得更快，那么它就可以满足大量的需求。NumPy、本地库、云计算环境等可以为其提供计算能力以启用训练过程。

但大多数组织不运行 Python，因此你可能需要将模型序列化为另一种格式，以便将其加载到 C/C++、C#、Java 或 JavaScript 应用程序中。随着 Open Neural Network Exchange (ONNX) 等开放格式的出现，这变得越来越容易，TensorFlow 等许多框架支持这种形式的混合使用。

第一步是训练模型。第二步是推理。直到几年前，TensorFlow 还要求你使用 Python 进行训练，但你可以将模型保存到磁盘并将它们加载到其他环境中。现在，你可以选择使用 Python、JavaScript 或 Swift 进行训练。

我们将讨论 TensorFlow.js，它被设计为在 JavaScript 浏览器中提供机器学习方法。使用 JavaScript 进行机器学习（尤其是在浏览器中）比你想象的更有意义。然而，为了理解机器学习与 JavaScript 和 WebAssembly 之间的关系，我们必须首先快速讨论一下硬件是如何改变我们的行业的。

硬件设施

程序员一般会认为他们的工作是编写在包含中央处理器 (CPU)、主存储器、存储系统、显示器和一些输入设备的计算机上运行的软件。虽然这仍然涵盖了很多软件开发，但我们拥有比大多数人意识到的更丰富的计算生态系统。计算机、平板电脑、游戏机、网络设备、智能手机、手表、嵌入式系统、物联网 (IoT) 设备和单片机 (SOC) 提供了一组执行软件的硬件系统。

著名的 Herb Sutter 写了一篇很棒的文章，叫做“欢迎来到丛林”，值得一读¹。在其中，他强调了过去几十年硬件发展对软件行业的影响。大约 30 年来，摩尔定律将更高的密度转化为更快的芯片²。如果你的软件运行缓慢，你只需等待 18 到 24 个月，它就会变得更快。

一旦我们不再能够制造更快、更复杂的芯片，我们就会利用额外的密度来制造更简单的芯片。多核系统成为常态。与以前的“免费午餐”不同，今天的开发人员必须编写疯狂的并发代码才能从额外的处理能力中获益，而且并不是每个问题都适合。另外，棘手的代码容易出错。这在很大程度上将我们推向函数式编程语言和不可变数据结构。

其他发展包括弹性云计算的出现、低延迟边缘计算的地理分布以及异构计算环境。这包括图形处理单元 (GPU)、现场可编程门阵列 (FPGA) 和专用集成电路 (ASIC)。

所有这些都与计算需要时间和电力的想法有关。成功的 IT 战略的很大一部分将是关于最小化时间成本、电力成本和延迟成本。这将影响事情的运作方式。将大量数据推送到云端以进行弹性、突发性训练是有意义的。但是，如果这导致大型模型，则由于规模问题，这些模型不太容易跨桌面/移动设备分布。同时，我们不希望互联网传感器和汽车刹车在关键时刻进行云调用。

我们正面临绝对的数据爆炸，这将使所有这些变得更加重要。为了在合理的时间内处理数据，我们需要访问硬件加速。这使数学并行化，因此它是可操作的。我们终于可以在浏览器中探索基于 JavaScript 的机器学习的概念。

Playground

TensorFlow Playground 是一个实验环境，非专家也可以通过直接操作直观地了解神经网络的工作原理³。这些模型越来越多地驱动真正感兴趣的系统，但理解它们所需的背景知识超出了非学术研究人员的能力范围。通过使用视觉表示，非技术用户可以获得直观的感受。

问题是，JavaScript 环境是单线程的，不能很好地支持训练神经网络所需的那种数学运算。Playground 环境允许你通过简单的用户界面操作更改训练过程中的所有参数。虽然改变这些参数的实验既简单又有趣，但它们会对结果的质量产生重大影响。为了快速实时地查看效果，运行时需要重新计算代码的重要部分；否则，直接操纵的意义就没有了。幸运的是，它们能够依靠浏览器中的 WebGL 来获得足够的性能来完成这项工作。我们稍后会进一步讨论这个想法。

TensorFlow.js

Playground 的成功让人们认识到将深度学习系统引入浏览器的必要性和可能性。在手机或平板电脑等潜在的低端设备上运行这种计算密集型系统，你可能会觉得有点奇怪。即使只是在强大的台式机上的浏览器中运行，也是一个稍微奇怪的概念。不过，这个想法还是有不少好处的。

点击链接下载一个基于深度学习的应用程序的体验，要比安装潜在的大量所需库容易得多。加载一个网页是零安装。如果这么容易，分享研究和实际应用就更容易了。这拓宽了深度学习研究人员之间的互动潜力，并使采用这些能力的应用程序更容易针对终端用户。

考虑到 JavaScript 的流行，要求 Web 开发人员学习用于机器学习的 Python 有点牵强。有大量的开源 JavaScript 代码用于构建各种形式的软件系统。能够利用这一切对机器学习系统的开发人员来说将是另一个好处。

手机和平板电脑已不仅仅是移动便携式电子设备。它们已成为诊断工具、银行系统的一部分、识别方法等等。人们正在编写新软件，通过分析网络摄像头录像来检测痴呆症的发作或中风的存在。与尝试从设备上获取数据相比，能够将应用程序推送到设备将更少受到监管负担的困扰，无论是用户体验质量还是诊断结果的隐私法合规性。

最后，许多可以运行这些下载的应用程序的设备都配备了强大而复杂的 GPU 供其使用。不仅可以在浏览器中执行深度学习系统，而且它实际上可能表现良好。

那么，我们为什么要讨论这一切？

TensorFlow.js 框架设计优雅，具有适用于各种设备的简洁 API。设计人员没有将它设计的大而全，而是选择创建一个可插入的后端来覆盖最广泛的系统。

后端的基础版本是基于 CPU 的 JavaScript 实现，可以在任何地方运行。所有代码都直接在 CPU 上执行，无需借助高级矢量扩展 (AVX)⁴ 等优化指令集。它基本上做到的可以在任何地方运行。你可能有更好的选择，但这是一个很好的后备方案。

下一个主要后端由 WebGL 加速。没有直接支持从 JavaScript 访问 GPU，而是通过几乎所有浏览器都支持的 WebGL。借助 3D 图形，开发人员可以在 GPU 上执行计算并将结果写入纹理图（texture）。这导致了一种快速方便的实现，几乎可以在所有主要的现代浏览器上运行良好。这类似于前面提到的 TensorFlow Playground 应用程序的情况。

有深度学习需求的服务器可以在 Node.js 中实现，所以下一个后端被设计成在比浏览器更自由的环境中运行。在 Node.js 中运行的应用程序可以从文件系统读取和写入，加载和使用原生库，并直接与普通的原生 TensorFlow 库通信。这些可以利用多核系统、GPU 或其他硬件加速设备，例如 Google 的张量处理单元 (TPU)。

你可以看到 API 设计的基本结构，如图 9-1 所示。

这里有很多东西要解读。首先，低级操作符通过上面显示的 Ops API 公开。需要在此级别深入了解详细信息的开发人员可以选择利用 API。然而，最重要的是，受 Keras 启发的 Layers API 在设计时充分考虑了开发人员的体验。这是一个简单的 API，用于使用简化的堆叠层方法定义复杂的神经网络架构。它隐藏了许多细节，但仍然通过相对简单的机制暴露出强大的行为。

无论客户端代码在 TensorFlow.js 中使用哪一层 API，针对这些 API 编写的应用程序都可以在后端覆盖的所有环境中移植。这包括浏览器内、浏览器外、有无硬件加速。这已经是一个了不起的设计成果，但是当代码能够利用硬件加速的优势时，在不同的环境中也能很好地运行。正如我们所见，这可以有多种形式，但我们不是针对最低公分母环境编写的，应用程序代码不需要绑定到一堆功能——测试代码以查看环境提供的内容。

使用 MobileNet 进行单次推理的结果，一百次运行的平均值如表 9-1 ⁵ 所示。此数据取自介绍 TensorFlow.js 的论文 TensorFlow.js：Web 及其他机器学习⁶。

作为一种趋势，结果并不令人惊讶，但细节却令人惊讶。在没有硬件加速的情况下，在原生 JavaScript 后端运行需要 3.5 秒。只需将后端更换为支持 WebGL 优化的后端，即使是非常适中的集成 GPU，在浏览器中的运行速度也会提高 71 倍。在具有更强大 GPU 的桌面上，我们实现了近 700 倍的加速。服务器应用程序在没有 GPU 的情况下运行，但在具有优化 AVX 指令的 CPU 的帮助下速度提高了近 40 倍。同样的环境，在更强大的 GPU 的支持下，产生了超过 1000 倍的惊人加速。请记住，应用程序没有任何变化，只是它运行的环境发生了变化。

这才是重点。从“免费午餐”时代到多核时代的转变迫使开发人员进行重大的编程更改以利用这些额外的资源。如果我们不针对自定义加速、多 CPU 的存在、云托管环境以及在浏览器中运行进行调优，我们的深度学习应用程序将成为代码无法维护的无底洞。相反，我们看到的是一个利用硬件加速的选项。这将我们带到我们将考虑的最后一个后端。

WebAssembly 后端

在 TensorFlow.js 发布后不久，该团队发布了一个新的 WebAssembly 后端。考虑到我们刚才讨论的范围，他们觉得有必要这样做，这可能会让你感到惊讶。但是，现在我想让你意识到这可能并不意味着从头开始。显然，还有一些工作要做，但其中相当一部分是使用现有代码完成的。特别是，扩展了 XNNPack 库以支持 WebAssembly 构建 ⁷。

这扩展了我们的后端集以包括图 9-2 中所示的后端。

这个例子突出了 WebAssembly 的一个突出用例。它并不是要取代 JavaScript，至少不是在所有情况下。它是关于扩展浏览器的可能性，而不必等待浏览器供应商之间的共识。WebGPU ⁸ 或类似的东西正在经历标准过程，但这可能需要几年时间。现有的用 C++ 编写并设计为跨多个平台进行优化和移植的库可以将其大部分功能带到今天的浏览器中，而无需等待。

WebAssembly 后端扩展了计算运行时，包括没有强大 GPU 的旧设备。可以优化代码以在广泛的平台上仍然运行良好。此外，WebAssembly 的设计者正在采用我们将在后面的章节中讨论的两个高级功能，即使用单指令多数据 (SIMD) 的并行化和多线程的使用 ⁹。

在图 9-3 中，我们对 MobileNet 的使用与之前类似。很明显，WebAssembly 后端本身的性能不如 WebGL 支持的后端。我们鼓励你使用 Wasm 后端而不是普通的 JavaScript 后端，因为它在任何地方都会表现得更好。根据模型的大小，Wasm 后端比 WebGL 后端更有意义，因为 WebGL 执行有一些固定成本。请注意，当将线程和 SIMD 并行化添加到 Wasm 组合中时，支持它们的平台将获得巨大的性能改进。

从图 9-4 中我们看到一个参数数量较少的不同模型的结果。在这种情况下，Wasm 的后端更有意义。请注意，Pixel 4 的性能相当，而 Linux 和 Mac 笔记本的性能则大幅提高。同样，线程和 SIMD 的使用极大地改善了基于 WebGL 的后端性能。

至少就目前而言，WebAssembly 并不总是在所有情况下都是最快的解决方案。这是一个相对较新的平台，所有潜在的优化都需要一段时间才能推出。然而，鉴于其年轻，WebAssembly 已经成为广泛计算环境中的游戏规则改变者。

本章的主要收获是 WebAssembly 及其相关技术可以为我们提供更多选择，在许多情况下是更好的选择，用于将高性能软件系统部署到不同的硬件和软件系统。随着时间的推移，我们会看到许多其他用途，但现在是扩大我们在语言和 WebAssembly 方面的选择的时候了。

注释