大幅提升 Web AI 模型測試效率：WebGPU、WebGL 和 Headless Chrome

Jason Mayes

François Beaufort

好消息！您已建立一個很酷的 Web AI 應用程式，可直接在使用者的裝置上執行機器學習模型。完全在用戶端網頁瀏覽器上執行，不必仰賴雲端。這項裝置端設計可強化使用者隱私權、提升效能，並大幅降低成本。

不過，這有個難題。TensorFlow.js 模型可在 CPU (WebAssembly) 和功能更強大的 GPU (透過 WebGL 和 WebGPU) 上運作。問題是：如何使用所選硬體，持續自動執行瀏覽器測試？

在為實際使用者部署機器學習模型前，您必須不斷迭代及改善模型，並維持一致性，才能比較機器學習模型的效能。

設定 GPU 的一致測試環境可能比預期更困難。在本篇文章中，我們將分享遇到的問題和解決方法，協助您改善應用程式效能。

這項功能不僅適用於 Web AI 開發人員！如果您是網頁遊戲或圖像的開發人員，這篇文章也能提供寶貴意見。

自動化工具箱中的工具

我們使用的工具如下：

• 環境：以 Linux 為基礎的 Google Colab notebook，已連線至 NVIDIA T4 或 V100 GPU。您可以使用其他雲端平台，例如 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，您應該會看到以下結果：

這不是個好的開始。很明顯，硬體偵測失敗。WebGL、WebGL2 和 WebGPU 基本上都已停用或僅限軟體。我們並非唯一遇到這個問題的人，許多人也遇到類似情況，並在線上討論 (包括 Chrome 官方支援管道：1、2)。

啟用 WebGPU 和 WebGL 支援

根據預設，Headless Chrome 會停用 GPU。如要在 Linux 上啟用此功能，請在啟動無頭 Chrome 時套用下列所有標記：

• --no-sandbox 標記會停用 Chrome 的安全沙箱，將瀏覽器程序與系統的其他部分隔離。不支援在沒有這個沙箱的情況下，以 root 權限執行 Chrome。
• --headless=new 標記會以全新且更強大的無頭模式執行 Chrome，且不會顯示任何 UI。
• --use-angle=vulkan 標記會指示 Chrome 為 ANGLE 使用 Vulkan 後端，將 OpenGL ES 2/3 呼叫轉譯為 Vulkan API 呼叫。
• --enable-features=Vulkan 標記可啟用 Vulkan 圖形後端，用於在 Chrome 中進行合成和算繪。
• --disable-vulkan-surface 旗標會停用 VK_KHR_surface vulkan 例項擴充功能。請改用 Bit blit 在畫面上呈現算繪結果，而非使用 Swapchain。
• --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 無頭版本，但這項版本將在某個時間點淘汰，改用新版。我們需要能與新版無頭 Chrome 搭配使用的解決方案，以便因應未來需求。

我們執行用於圖像辨識的 TensorFlow.js 範例網頁，確認 GPU 的使用率偏低，並在該網頁中訓練辨識服裝樣本的模型 (有點像機器學習的「Hello World」)。

在一般機器上，50 個訓練週期 (稱為「epoch」) 的執行時間應各自少於 1 秒。在預設狀態下呼叫無頭 Chrome，我們可以將 JavaScript 資訊主機輸出記錄到 Node.js 伺服器端指令列，查看這些訓練週期實際上需要多長時間。

如預期，每個訓練週期都比預期時間長得多 (幾秒)，這表示 Chrome 已改回使用純粹的 JS CPU 執行作業，而非利用 GPU：

修正驅動程式並使用適當的標記組合後，重新執行 TensorFlow.js 訓練範例，即可大幅加快訓練週期。

摘要

自 2017 年推出以來，Web AI 的使用量呈現指數級成長。透過 WebGPU、WebGL 和 WebAssembly 等瀏覽器技術，可進一步加快用戶端的機器學習模型數學運算。

截至 2023 年，TensorFlow.js 和 MediaPipe Web 的模型和程式庫下載次數已超過 10 億次，這項成就不僅是歷史性的里程碑，也顯示出網頁開發人員和工程師正轉向採用 AI 來打造新一代網頁應用程式，進而開發出令人驚豔的解決方案。

隨著使用量大幅增加，責任也隨之加重。在實際工作系統中，使用者需要在真實的瀏覽器環境中測試用戶端的瀏覽器 AI 模型，同時還要能夠擴充、自動化，並在已知的標準化硬體設定中執行。

只要善用新版無頭 Chrome 和 Puppeteer 的強大功能，您就能在標準化且可複製的環境中放心測試這類工作負載，確保獲得一致且可靠的結果。

總結

我們的文件中提供逐步指南，您可以自行嘗試完成設定。

如果您覺得這篇文章實用，歡迎在 LinkedIn、X (舊稱 Twitter) 或您使用的任何社群媒體上，使用主題標記 #WebAI 來推薦這篇文章。歡迎提供任何意見，讓我們知道日後該如何撰寫這類內容。

在 GitHub 存放區中新增星號，即可接收日後的任何更新。

特別銘謝

在此感謝 Chrome 團隊的所有成員，協助我們排解這個解決方案中遇到的驅動程式和 WebGPU 問題，特別感謝 Jecelyn Yeen 和 Alexandra White 協助撰寫這篇網誌文章。感謝 Yuly Novikov、Andrey Kosyakov 和 Alex Rudenko，他們在打造最終可行解決方案的過程中發揮了重要作用。