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

Jason Mayes

François Beaufort

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

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

在迭代和改进机器学习模型后，您需要先确保一致性，然后再将其部署供真实用户在其设备上使用，这样才能比较机器学习模型在一段时间内的性能。

使用 GPU 设置一致的测试环境可能比预期更难。在本博文中，我们将分享我们遇到的问题以及解决方法，以便您提高应用的性能。

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

自动化工具箱中包含哪些内容

我们使用的是以下工具：

• 环境：连接到 NVIDIA T4 或 V100 GPU 的基于 Linux 的 Google Colab Notebook。您也可以使用其他云平台（例如 Google Cloud [GCP]），如果您愿意的话。
• 浏览器：Chrome 支持 WebGPU，这是一种强大的 WebGL 后继者，可将现代 GPU API 的进步引入到 Web 中。
• 自动化：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，您应该会看到以下结果：

开局不顺利。很明显，硬件检测失败了。 WebGL、WebGL2 和 WebGPU 实际上处于停用状态或仅限软件。我们并不是唯一遇到此问题的人，网上有许多讨论，其中包括 Chrome 官方支持渠道 (1)、(2) 上，都有遇到类似问题的用户。

启用 WebGPU 和 WebGL 支持

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

• --no-sandbox 标志会停用 Chrome 的安全沙盒，该沙盒会将浏览器进程与系统的其余部分隔离开来。不支持在没有此沙盒的情况下以 root 身份运行 Chrome。
• --headless=new 标志可让 Chrome 以经过改进的新版无头模式运行，而无需任何可见的界面。
• --use-angle=vulkan 标志会指示 Chrome 为 ANGLE 使用 Vulkan 后端，后者会将 OpenGL ES 2/3 调用转换为 Vulkan API 调用。
• --enable-features=Vulkan 标志用于在 Chrome 中启用 Vulkan 图形后端以进行合成和光栅化。
• --disable-vulkan-surface 标志会停用 VK_KHR_surface Vulkan 实例扩展。系统使用位复制（而非使用交换链）在屏幕上呈现渲染结果。
• --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 问题，且值会从“已停用”更改为“仅限软件”。

不过，硬件加速功能仍不可用，系统未检测到 NVIDIA T4 GPU。

安装正确的 GPU 驱动程序

我们与 Chrome 团队的部分 GPU 专家一起，更仔细地研究了 chrome://gpu 的输出。我们发现 Linux Colab 实例上安装的默认驱动程序存在问题，导致 Vulkan 出现问题，进而导致 Chrome 无法在 GL_RENDERER 级别检测 NVIDIA T4 GPU，如以下输出所示。这会导致 Headless Chrome 出现问题。

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

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

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

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

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

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

幕后揭秘：我们团队的调查

经过大量研究，我们没有找到适用于我们需要在 Google Colab 中执行的环境的有效方法，但有一些有希望的帖子适用于其他环境，这让我们看到了希望。最终，我们无法在 Colab NVIDIA T4 环境中复制他们的成功，因为我们遇到了 2 个主要问题：

1. 某些标志组合允许检测 GPU，但不允许您实际使用 GPU。
2. 第三方提供的有效解决方案示例使用了旧版无头 Chrome，该版本将在某个时间点被弃用，取而代之的是新版本。我们需要一个可与新版 Headless Chrome 搭配使用的解决方案，以便更好地应对未来需求。

我们通过运行用于图片识别的示例 TensorFlow.js 网页来确认 GPU 利用率偏低，其中我们训练了一个模型来识别服装样本（有点像机器学习的“Hello World”）。

在普通机器上，50 个训练周期（称为“时代”）的运行时间应均不超过 1 秒。在默认状态下调用无头 Chrome 时，我们可以将 JavaScript 控制台输出记录到 Node.js 服务器端命令行，以了解这些训练周期的实际速度。

正如预期，每个训练周期都比预期要长得多（几秒钟），这表明 Chrome 已回退到普通旧版 JS CPU 执行，而不是利用 GPU：

修复驱动程序并为无头 Chrome 使用正确的标志组合后，重新运行 TensorFlow.js 训练示例会使训练周期变得更快。

摘要

自 2017 年创建以来，Web AI 呈指数级增长。借助 WebGPU、WebGL 和 WebAssembly 等浏览器技术，机器学习模型的数学运算可以在客户端进一步加速。

截至 2023 年，TensorFlow.js 和 MediaPipe Web 的模型和库下载量已超过 100 亿次，这是一个历史性里程碑，表明 Web 开发者和工程师正在转变，在下一代 Web 应用中采用 AI 技术，打造一些真正令人惊叹的解决方案。

随着使用量不断增加，责任也越来越大。在生产系统中达到这种使用水平后，就需要在真实的浏览器环境中测试基于浏览器的客户端 AI 模型，同时还需要可扩缩、可自动化，并且在已知的标准化硬件设置中进行测试。

通过利用新版无头 Chrome 和 Puppeteer 的强大组合，您可以放心地在标准化且可重现的环境中测试此类工作负载，确保获得一致且可靠的结果。

小结

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

如果您觉得本指南有用，欢迎在 LinkedIn、X（以前称为 Twitter） 或您使用的任何社交网络上，使用 #WebAI 标签进行宣传。我们非常期待收到您的反馈，以便日后创作更多类似内容。

在 GitHub 代码库中为项目添加星标，以接收日后发布的所有更新。

致谢

非常感谢 Chrome 团队中的所有成员帮助我们调试在此解决方案中遇到的驱动程序和 WebGPU 问题，特别感谢 Jecelyn Yeen 和 Alexandra White 帮助我们润色这篇博文。感谢 Yuly Novikov、Andrey Kosyakov 和 Alex Rudenko，他们在打造最终可行的解决方案方面发挥了重要作用。