强化 Web AI 模型测试：WebGPU、WebGL 和 Headless Chrome

贾森·梅耶斯
Jason Mayes
Twitter GitHub
弗朗索瓦·博福
François Beaufort
GitHub

好消息！您已构建了一个超酷的 Web AI 应用，该应用可直接在用户设备上运行机器学习模型。它完全在客户端网络浏览器上运行，无需依赖云端。这种设备端设计可显著增强用户隐私保护、提高性能并降低费用。

但是，这其中存在一个障碍。您的 TensorFlow.js 模型可以在 CPU (WebAssembly) 和更强大的 GPU（通过 WebGL 和 WebGPU）上运行。问题是：如何使用所选硬件以一致的方式自动执行浏览器测试？

在为实际用户在其设备上使用之前，保持一致性对于比较您迭代和改进机器学习模型在一段时间内的性能来说至关重要。

使用 GPU 设置一致的测试环境可能比预期更难。在这篇博文中，我们将分享我们所面临的问题以及我们是如何解决这些问题的，以帮助您提升应用的性能。

这不仅仅适用于 Web AI 开发者！如果您从事网页游戏或图形方面的工作，此博文对您来说也很有价值。

自动化工具箱中的内容

我们使用的方法如下：

  • 环境：连接到 NVIDIA T4 或 V100 GPU 的基于 Linux 的 Google Colab 笔记本。如果您愿意，可以使用 Google Cloud (GCP) 等其他云平台。
  • 浏览器：Chrome 支持 WebGPU，这是 WebGL 的强大后继产品，推动了现代 GPU API 的进步。
  • 自动化操作：Puppeteer 是一个 Node.js 库，可让您使用 JavaScript 以编程方式控制浏览器。借助 Puppeteer，我们可以在无头模式下自动执行 Chrome，这意味着浏览器在服务器上没有可见界面的情况下运行。我们使用的是经过改进的新无头模式，而不是旧版形式。

验证环境

若要检查 Chrome 中是否已开启硬件加速，最佳方法是在地址栏中输入 chrome://gpu。您可以通过 console.log 以编程方式使用 Puppeteer 执行等效操作，也可以将完整报告保存为 PDF 以手动检查：

/* Incomplete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters: Expands later
const browser = await puppeteer.launch({
  headless: 'new',
  args:  ['--no-sandbox']
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({ path: './gpu.pdf' });

await browser.close();

打开 chrome://gpu，您应该会看到以下结果：

图形功能状态
OpenGL: 已停用
Vulkan： 已停用
WebGL： 仅支持软件，不支持硬件加速。
WebGL2： 仅支持软件，不支持硬件加速。
WebGPU： 已停用

检测到问题。
已通过屏蔽名单或命令行停用 WebGPU。

良好的开端。很显然，硬件检测失败了。WebGL、WebGL2 和 WebGPU 基本上已停用，或者仅支持软件。这一问题不止我们一个人，网上有很多人讨论过类似的情况，包括官方 Chrome 支持渠道 (1) (2)。

启用 WebGPU 和 WebGL 支持

默认情况下，无头 Chrome 会停用 GPU。如需在 Linux 上启用 Headless Chrome，请在启动无头 Chrome 时应用以下所有标志：

  • --no-sandbox 标志会停用 Chrome 的安全沙盒，此沙盒会将浏览器进程与系统的其余部分隔离开来。不支持在不使用此沙盒的情况下以 root 用户身份运行 Chrome。
  • --headless=new 标志：通过经过改进的全新无头模式运行 Chrome，无任何可见界面。
  • --use-angle=vulkan 标志会告知 Chrome 为 ANGLE 使用 Vulkan 后端，这会将 OpenGL ES 2/3 调用转换为 Vulkan API 调用。
  • --enable-features=Vulkan 标志用于启用 Vulkan 图形后端，以便在 Chrome 中进行合成和光栅化。
  • --disable-vulkan-surface 标志用于停用 VK_KHR_surface vulkan 实例扩展程序。它不使用交换链，而是针对屏幕上的当前渲染结果使用 Bit blit。
  • --enable-unsafe-webgpu 标志用于在 Linux 上的 Chrome 中启用实验性 WebGPU API，并停用适配器屏蔽名单。

现在，我们将到目前为止所做的所有更改合并在一起。以下是完整的脚本。

/* Complete example.js */
import puppeteer from 'puppeteer';

// Configure launch parameters
const browser = await puppeteer.launch({
  headless: 'new',
  args: [
    '--no-sandbox',
    '--headless=new',
    '--use-angle=vulkan',
    '--enable-features=Vulkan',
    '--disable-vulkan-surface',
    '--enable-unsafe-webgpu',
  ]
});

const page = await browser.newPage();
await page.goto('chrome://gpu');

// Verify: log the WebGPU status or save the GPU report as PDF
const txt = await page.waitForSelector('text/WebGPU');
const status = await txt.evaluate(g => g.parentElement.textContent);
console.log(status);
await page.pdf({path: './gpu.pdf'});

await browser.close();

再次运行该脚本。未检测到 WebGPU 问题，并且值从“已停用”更改为“仅限软件”。

图形功能状态
OpenGL: 已停用
Vulkan： 已停用
WebGL： 仅支持软件，不支持硬件加速。
WebGL2： 仅支持软件，不支持硬件加速。
WebGPU： 仅支持软件，不支持硬件加速。

但是，硬件加速仍然不可用，因此系统未检测到 NVIDIA T4 GPU。

安装正确的 GPU 驱动程序

我们请 Chrome 团队的一些 GPU 专家深入调查了 chrome://gpu 的输出。我们发现 Linux Colab 实例上安装的默认驱动程序存在问题，从而导致 Vulkan 出现问题，进而导致 Chrome 无法在 GL_RENDERER 级别检测到 NVIDIA T4 GPU，如以下输出所示。这会导致无头 Chrome 出现问题。

默认输出不检测 NVIDIA T4 GPU。
驾驶员信息
GL_RENDERER ANGLE（Google，Vulkan 1.3.0 (SwiftShader Device (Subzero) (0x0000C0DE))，SwiftShader 驱动程序-5.0.0

安装兼容的正确驱动程序即可解决此问题。

。
安装驱动程序后更新了输出
驾驶员信息
GL_RENDERER ANGLE（NVIDIA Corporation、Tesla T4/PCIe/SSE2、OpenGL ES 3.2 NVIDIA 525.105.17）

如需安装正确的驱动程序，请在设置期间运行以下命令。最后两行可帮助您记录 NVIDIA 驱动程序检测到的输出以及 vulkaninfo。

apt-get install -y vulkan-tools libnvidia-gl-525

// Verify the NVIDIA drivers detects along with vulkaninfo
nvidia-smi
vulkaninfo --summary

现在再次运行该脚本，我们得到以下结果。🎉

图形功能状态
OpenGL: 已启用
Vulkan： 已启用
WebGL： 已硬件加速，但性能下降。
WebGL2： 已硬件加速，但性能下降。
WebGPU： 已硬件加速，但性能下降。

通过在运行 Chrome 时使用正确的驱动程序和标志，我们现在可通过全新的无头模式支持 WebGPU 和 WebGL。

幕后：我们的团队调查

经过大量研究后，我们没有找到可在 Google Colab 中执行的环境的有效方法，但有一些乐于发现的博文在其他环境中也能成功运行，前景可观。最终，我们无法在 Colab NVIDIA T4 环境中复制他们的成功，因为我们有两个主要问题：

  1. 某些标志组合允许检测 GPU，但不允许您实际使用 GPU。
  2. 第三方的有效解决方案示例使用的是旧版 Chrome 无头版本，此版本日后将被弃用，取而代之的是新版本。我们需要一个可与新版无头 Chrome 配合使用的解决方案，以便更好地适应未来变化。

我们通过运行用于图像识别的 TensorFlow.js 网页示例来确认 GPU 利用率过低，并由此训练了一个能够识别服装样本的模型（有点像机器学习的“你好世界”）。

在常规机器上，50 个训练周期（称为周期）的运行时间应少于 1 秒。调用处于默认状态的 Headless Chrome，我们可以将 JavaScript 控制台输出记录到 Node.js 服务器端命令行，以查看这些训练周期的实际速度。

不出所料，每个训练周期花费的时间都比预期长得多（几秒），这表明 Chrome 已回退到普通的旧 JS CPU 执行，而不是使用 GPU：

训练周期以较慢的频率移动。
图 1：显示每个训练周期执行时长（秒）的实时捕获。

在修复驱动程序并为无头 Chrome 使用正确的标志组合后，重新运行 TensorFlow.js 训练示例可以大幅加快训练周期。

周期的速度有所提高...
图 2：显示周期速度的实时拍摄。

摘要

自 2017 年问世以来，Web AI 一直在快速增长。利用 WebGPU、WebGL 和 WebAssembly 等浏览器技术，可以在客户端进一步加速机器学习模型的数学运算。

截至 2023 年，TensorFlow.js 和 MediaPipe Web 的模型和库下载量已超过 10 亿次，这标志着 Web 开发者和工程师正在转变发展方向，在新一代 Web 应用中融入 AI 技术，从而打造出一些真正不可思议的解决方案。

成功的使用伴随着很大的责任。在生产系统的这种使用水平下，需要在真实的浏览器环境中测试基于浏览器的客户端 AI 模型，同时还需要在已知的标准化硬件设置中实现可伸缩、可自动化和。

通过利用新的 Headless Chrome 和 Puppeteer 的强大功能，您可以在标准化的可复制环境中放心地测试此类工作负载，确保获得一致且可靠的结果。

小结

我们的文档中提供了分步指南，因此您可以自行尝试完成整个设置。

如果您觉得这有用，可以在 LinkedIn、X（以前称为 Twitter）或您使用 #WebAI 标签 #WebAI 的任何社交网络上公开致谢。我们非常期待收到您的反馈，以便我们知道将来可以创作更多类似的文章。

在 GitHub 代码库中添加星标，以接收将来的任何更新。

致谢

衷心感谢 Chrome 团队中帮助调试此解决方案中遇到的驱动程序和 WebGPU 问题的每个人，并特别感谢 Jecelyn Yeen 和 Alexandra White 对这篇博文的帮助。感谢 Yuly Novikov、Andrey Kosyakov 和 Alex Rudenko 为创建最终有效的解决方案贡献了力量。