大幅提升 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。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 時套用下列所有標記：

  • --no-sandbox 旗標會停用 Chrome 的安全沙箱，將瀏覽器程序與系統的其他部分隔離。系統不支援在沒有這個沙箱的情況下以根權限執行 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 裝置 (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 環境中複製這些成功經驗，因為有 2 個重要問題：

  1. 部分標記組合可讓您偵測 GPU，但不允許您實際使用 GPU。
  2. 由第三方提供的有效解決方案範例，其中使用的是舊版 Chrome 無頭版本，該版本日後將會淘汰，並改用新版本。我們需要能和全新的 Headless Chrome 搭配使用的解決方案，才能更有效地因應未來趨勢。

我們執行圖片辨識範例的 TensorFlow.js 網頁，確認 GPU 使用率偏低，並訓練了一個模型來辨識服裝樣本 (像是機器學習的「Hello World」)。

在一般機器上，50 個訓練週期 (稱為 Epoch 紀元) 每次的執行時間應少於 1 秒。呼叫無頭 Chrome 做為預設狀態時，我們可以將 JavaScript 控制台輸出內容記錄至 Node.js 伺服器端指令列，瞭解這些訓練週期實際花費的速度。

一如預期，每個訓練週期所花費的時間超過預期 (數秒)，這表示 Chrome 已不再使用簡單的 JS CPU 執行，而非使用 GPU：

訓練週期會較慢。
圖 1：即時擷取顯示每個訓練週期所需的時間 (秒)。

修正驅動程式並使用 Headless Chrome 的正確標記組合後，重新執行 TensorFlow.js 訓練範例能夠大幅加快訓練週期。

訓練週期的速度加快。
圖 2：即時擷取顯示訓練週期的速度。

摘要

自 2017 年以來，Web AI 的成長速度。透過 WebGPU、WebGL 和 WebAssembly 等瀏覽器技術，機器學習模型的數學作業可以在用戶端加速。

截至 2023 年，TensorFlow.js 和 MediaPipe Web 跨超過 10 億次的模型和程式庫下載次數超過了 10 億次。這項歷史里程碑，也意味著網頁開發人員和工程師如何改用新一代網頁應用程式中的 AI 技術，打造出真正令人驚豔的解決方案。

隨著使用者持續獲得良好成效，我們必須負起重大責任。因此，在真正的瀏覽器環境中測試用戶端、瀏覽器式 AI 模型時，也需要在真正的瀏覽器環境中測試用戶端、瀏覽器式 AI 模型，而且要在已知的標準化硬體設定中都能擴充。

透過結合全新無頭 Chrome 和 Puppeteer 的強大功能，您可以在標準化且可複製的環境中放心測試這類工作負載，以確保結果一致且可靠。

總結

我們的說明文件提供逐步指南，可讓您自行嘗試完成設定程序。

如果覺得這很實用，請在 LinkedIn、X (前身為 Twitter) 或任何您使用的社群網路中，使用 #WebAI 主題標記進行宣傳。希望您能與我們分享寶貴意見，讓我們日後能據此撰寫更多這類內容。

在 GitHub 存放區上加上星號，以接收日後的任何更新。

特別銘謝

非常感謝 Chrome 團隊協助偵錯此解決方案中遇到的驅動程式和 WebGPU 問題，特別要感謝 Jecelyn Yeen 和 Alexandra White 撰寫這篇網誌文章。感謝 Yuly Novikov、Andrey Kosyakov 和 Alex Rudenko 這名人士，在建立最終的可運作解決方案方面，