在 Google Colab 中进行 Web AI 模型测试

Jason Mayes
François Beaufort
François Beaufort
GitHub

使用 GPU 设置一致的测试环境可能比预期更难。以下是测试基于浏览器的客户端 AI 模型的步骤，可在真实的浏览器环境中进行测试，同时具有可伸缩性、可自动化，并且在已知的标准化硬件设置中进行测试。

在本例中，浏览器是具有硬件支持的真实 Chrome 浏览器，而不是软件模拟。

无论您是 Web AI、Web 游戏或图形开发者，还是对 Web AI 模型测试感兴趣，本指南都适合您。

第 1 步：创建新的 Google Colab 笔记本

1. 前往 colab.new 创建新的 Colab 笔记本。界面应如下图所示。 2. 按照提示登录您的 Google 账号。
新版 Colab 的屏幕截图
图 1：新的 Colab 笔记本。

第 2 步：连接到启用了 T4 GPU 的服务器

  1. 点击记事本右上角附近的关联图标 。
  2. 选择更改运行时类型：
    放大显示更改运行时的步骤的屏幕截图。
    图 2。在 Colab 界面中更改运行时。
  3. 在模态窗口中，选择 T4 GPU 作为硬件加速器。连接后，Colab 将使用连接有 NVIDIA T4 GPU 的 Linux 实例。
    “更改运行时类型”模块的屏幕截图。
    图 3：在“硬件加速器”下，选择 T4 GPU。
  4. 点击保存。
  5. 点击 Connect 按钮以连接到运行时。过了一段时间后，该按钮将显示一个绿色对勾标记，以及 RAM 和磁盘使用情况图表。这表示系统已成功使用所需的硬件创建服务器。

干得漂亮，您刚刚创建了一个挂接了 GPU 的服务器。

第 3 步：安装正确的驱动程序和依赖项

  1. 将以下两行代码复制并粘贴到笔记本的第一个代码单元中。在 Colab 环境中，命令行执行会在开头添加一个感叹号。

    !git clone https://github.com/jasonmayes/headless-chrome-nvidia-t4-gpu-support.git
!cd headless-chrome-nvidia-t4-gpu-support && chmod +x scriptyMcScriptFace.sh && ./scriptyMcScriptFace.sh

    # Update, install correct drivers, and remove the old ones.
apt-get install -y vulkan-tools libnvidia-gl-525

# Verify NVIDIA drivers can see the T4 GPU and that vulkan is working correctly.
nvidia-smi
vulkaninfo --summary

# Now install latest version of Node.js
npm install -g n
n lts
node --version
npm --version

# Next install Chrome stable
curl -fsSL https://dl.google.com/linux/linux_signing_key.pub | sudo gpg --dearmor -o /usr/share/keyrings/googlechrom-keyring.gpg
echo "deb [arch=amd64 signed-by=/usr/share/keyrings/googlechrom-keyring.gpg] http://dl.google.com/linux/chrome/deb/ stable main" | sudo tee /etc/apt/sources.list.d/google-chrome.list
sudo apt update
sudo apt install -y google-chrome-stable

# Start dbus to avoid warnings by Chrome later.
export DBUS_SESSION_BUS_ADDRESS="unix:path=/var/run/dbus/system_bus_socket"
/etc/init.d/dbus start

  2. 点击单元格旁边的 以执行代码。

    新版 Colab 的屏幕截图
    图 4。

  3. 代码执行完毕后，请验证 nvidia-smi 是否输出了类似于以下屏幕截图的内容，以确认您确实已连接 GPU 且服务器上已识别到该 GPU。您可能需要滚动到日志中的较早部分才能查看此输出。

    图 5：查找以“NVIDIA-SMI”开头的输出。

第 4 步：使用和自动化无头 Chrome

  1. 点击 Code（代码）按钮以添加新的代码单元。
  2. 然后，您可以编写自定义代码，以使用首选参数调用 Node.js 项目（或直接在命令行中调用 google-chrome-stable）。我们提供了以下两种示例。

第 A 部分：直接在命令行中使用无头 Chrome

# Directly call Chrome to dump a PDF of WebGPU testing page
# and store it in /content/gpu.pdf
!google-chrome-stable \
--no-sandbox \
--headless=new \
--use-angle=vulkan \
--enable-features=Vulkan \
--disable-vulkan-surface \
--enable-unsafe-webgpu \
--print-to-pdf=/content/gpu.pdf https://webgpureport.org

在此示例中，我们将生成的 PDF 截图存储在 /content/gpu.pdf 中。如需查看该文件，请展开内容 。然后，点击 将 PDF 文件下载到本地机器。

新版 Colab 的屏幕截图
图 6：请观看此 Colab 界面屏幕截图，了解下载 PDF 的步骤。

第 B 部分：使用 Puppeteer 控制 Chrome

我们提供了一个使用 Puppeteer 控制无头 Chrome 的极简示例，您可以按如下方式运行该示例：

# Call example node.js project to perform any task you want by passing
# a URL as a parameter
!node headless-chrome-nvidia-t4-gpu-support/examples/puppeteer/jPuppet.js chrome://gpu

在 jPuppet 示例中，我们可以调用 Node.js 脚本来创建屏幕截图。但具体是怎么运作的？请查看此 jPuppet.js 中 Node.js 代码的演示。

jPuppet.js Node 代码细分

首先，导入 Puppeteer。这样，您就可以使用 Node.js 远程控制 Chrome：

import puppeteer from 'puppeteer';

接下来，检查传递给 Node 应用的命令行参数。 确保设置第三个参数，该参数表示要导航到的网址。您需要在此处检查第三个参数，因为前两个参数会调用 Node 本身和我们正在运行的脚本。第 3 个元素实际上包含传递给 Node 程序的第 1 个参数：

const url = process.argv[2];
if (!url) {
  throw "Please provide a URL as the first argument";
}

现在，定义一个名为 runWebpage() 的异步函数。这会创建一个使用命令行参数配置的浏览器对象，以便以我们需要的方式运行 Chrome 二进制文件，从而使 WebGL 和 WebGPU 正常运行，如启用 WebGPU 和 WebGL 支持中所述。

async function runWebpage() {
  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();

然后，添加一个事件监听器，以便在网页执行 JavaScript 时监听 console.log 事件。这样，您就可以在 Node 命令行中记录消息，还可以检查控制台文本是否包含用于触发屏幕截图的特殊字词（在本例中为 captureAndEnd），然后在 Node 中结束浏览器进程。对于需要在截取屏幕截图之前执行一些工作且执行时间不确定的网页，此方法非常有用。

page.on('console', async function(msg) {
  console.log(msg.text());
  if (msg.text() === 'captureAndEnd') {
    await page.screenshot({ path: '/content/screenshotEnd.png' });
    await browser.close();
  }
});

最后，命令网页访问指定的网址，并在网页加载后截取初始屏幕截图。

如果您选择截取 chrome://gpu 的屏幕截图，则可以立即关闭浏览器会话，而无需等待任何控制台输出，因为此页面不受您自己的代码控制。

  await page.goto(url,  { waitUntil: 'networkidle2' });
  await page.screenshot({path: '/content/screenshot.png'});
  if (url === 'chrome://gpu') {
    await browser.close();
  }
}
runWebpage();

修改 package.json

您可能已经注意到，我们在 jPuppet.js 文件开头使用了 import 语句。您的 package.json 必须将类型值设置为 module，否则您会收到模块无效的错误。

 {
    "dependencies":  {
      "puppeteer": "*"
    },
    "name": "content",
    "version": "1.0.0",
    "main": "jPuppet.js",
    "devDependencies": {},
    "keywords": [],
    "type": "module",
    "description": "Node.js Puppeteer application to interface with headless Chrome with GPU support to capture screenshots and get console output from target webpage"
}

这就是全部内容！使用 Puppeteer 可以更轻松地以编程方式与 Chrome 交互。

成功

现在，我们可以验证 TensorFlow.js Fashion MNIST 分类器是否能够在浏览器中使用 GPU 进行客户端处理，从而正确识别图片中的一条裤子。

您可以将其用于任何基于 GPU 的客户端工作负载，从机器学习模型到图形和游戏测试。

新版 Colab 的屏幕截图
图 7：成功捕获了能够在浏览器中实时识别服装的客户端 GPU 加速型 TensorFlow.js 模型

资源

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