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

Jason Mayes

François Beaufort

好消息！您已构建了一个很酷的 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 支持

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

• --no-sandbox 标志可停用 Chrome 的安全沙盒，这会将浏览器进程与系统的其余部分隔离开来。在没有此沙盒的情况下，不支持以 root 身份运行 Chrome。
• --headless=new 标志使用经过改进的新无头模式运行 Chrome，但无任何可见界面。
• --use-angle=vulkan 标志会告知 Chrome 将 Vulkan 后端用于 ANGLE，这会将 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 公司、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 无头版本有时将被弃用，取而代之的是新版本。我们需要一种可与新的无头 Chrome 配合使用的解决方案，以便更好地适应未来变化。

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

在常规机器上，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 模型，同时还需要在已知的标准化硬件设置中测试可扩展、可自动化的模型。

当我们说“真正的浏览器环境”时，我们指的是这是一个具有硬件支持的真实 Chrome 浏览器，而不是软件模拟。

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

小结

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

如果您觉得这些内容有用，请在 LinkedIn、X（以前称为 Twitter） 或您使用的任何社交网络上发布（使用 # 标签 #WebAI）。非常期待收到您的任何反馈，以便我们了解以后会编写更多此类内容。

在 GitHub 代码库上加注星标，以接收任何后续更新。

致谢

非常感谢 Chrome 团队中协助调试此解决方案中遇到的驱动程序和 WebGPU 问题的每一位用户，并特别感谢 Jecelyn Yeen 和 Alexandra White 在这篇博文中提供帮助的每一个人。感谢 Yuly Noviov、Andrey Kosya 之后 [Andrey Kosya 之后] 和 Alex Rudenko 对制定最终有效解决方案的大力支持。