好消息!您已构建了一个很酷的 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: | 已停用 |
检测到问题。 |
开局并不好。很明显,硬件检测失败。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 出现问题。
驱动程序信息 | |
---|---|
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 环境中复制它们的成功,因为我们有两个关键问题:
- 某些标志组合允许检测 GPU,但无法实际使用 GPU。
- 第三方的有效解决方案示例使用的是旧版 Chrome 无头版本,旧版 Chrome 无头版本有时将被弃用,取而代之的是新版本。我们需要一种可与新的无头 Chrome 配合使用的解决方案,以便更好地适应未来变化。
我们通过运行用于图像识别的示例 TensorFlow.js 网页来确认 GPU 利用率过低,并借此训练一个模型来识别服装样本(有点像机器学习的“hello world”)。
在常规机器上,50 个训练周期(称为周期)的运行时间应少于 1 秒。在默认状态下调用 Headless Chrome,我们可以将 JavaScript 控制台输出记录到 Node.js 服务器端命令行,以查看这些训练周期的实际运行速度。
不出所料,每个训练周期所花的时间比预期要长得多(几秒),这表明 Chrome 已回退到普通的 JS CPU 执行方式,而不是利用 GPU:
在修复驱动程序并对无头 Chrome 使用正确的标志组合后,重新运行 TensorFlow.js 训练示例可以显著加快训练周期。
摘要
自 2017 年问世以来,Web AI 实现了指数级增长。借助 WebGPU、WebGL 和 WebAssembly 等浏览器技术,可以进一步在客户端加快机器学习模型的数学运算。
截至 2023 年,TensorFlow.js 和 MediaPipe Web 的模型和库下载量超过 10 亿,这是一个历史里程碑,标志着 Web 开发者和工程师如何开始在其下一代 Web 应用中运用 AI 技术,打造一些真正卓越的解决方案。
使用得当,责任越大。生产系统使用这种水平时,需要在真实的浏览器环境中测试基于浏览器的客户端 AI 模型,同时还需要在已知的标准化硬件设置中测试可扩展、可自动化的模型。
通过结合使用新的 Headless Chrome 和 Puppeteer 的强大功能,您可以在可复制的标准化环境中放心地测试此类工作负载,确保获得一致且可靠的结果。
小结
我们的文档中提供了分步指南,因此您可以自行尝试完整的设置过程。
如果您觉得这些内容有用,请在 LinkedIn、X(以前称为 Twitter) 或您使用的任何社交网络上发布(使用 # 标签 #WebAI)。非常期待收到您的任何反馈,以便我们了解以后会编写更多此类内容。
在 GitHub 代码库上加注星标,以接收任何后续更新。
致谢
非常感谢 Chrome 团队中协助调试此解决方案中遇到的驱动程序和 WebGPU 问题的每一位用户,并特别感谢 Jecelyn Yeen 和 Alexandra White 在这篇博文中提供帮助的每一个人。感谢 Yuly Noviov、Andrey Kosya 之后 [Andrey Kosya 之后] 和 Alex Rudenko 对制定最终有效解决方案的大力支持。