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

Jason Mayes

François Beaufort

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

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

將機器學習模型效能隨時間經過不斷疊代與改進時，保持一致性是非常重要的，因為在將模型部署到實際裝置之前，則維持一致性。

透過 GPU 設定一致的測試環境並不容易。我們將在這篇網誌文章中分享目前遇到的問題及解決方式，以便您改善應用程式效能。

這不只適用於網頁 AI 開發人員！如果你負責處理網路遊戲或圖片，這篇文章也十分實用。

自動化工具箱提供哪些功能

以下是我們使用的工具：

• 環境：以 Linux 為基礎的 Google Colab 筆記本，連接到 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 支援功能

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

• --no-sandbox 標記會停用 Chrome 安全沙箱，將瀏覽器程序與系統其他部分隔離。在沒有這個沙箱的情況下，系統將不支援以根權限執行 Chrome。
• --headless=new 標記會在沒有任何可見 UI 的情況下，使用全新改良的無頭模式執行 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 問題，且值會從停用變更為僅限軟體。

不過，還是無法使用硬體加速功能，系統也不會偵測 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 無頭版本的工作解決方案範例，在某個時間點將淘汰，並改用新版本。我們需要與新版 Headless Chrome 搭配運作的解決方案，進一步保障及因應未來需求。

我們執行了圖片辨識範例 TensorFlow.js 網頁，確認 GPU 使用率不足，並在其中訓練模型，辨識服飾樣本 (類似機器學習的「Hello World」)。

在一般機器中，每個訓練週期 (稱為訓練週期) 的執行週期應在 1 秒內不到 1 秒。以預設狀態呼叫 Headless Chrome 時，我們可以將 JavaScript 控制台輸出內容記錄到 Node.js 伺服器端指令列，查看實際執行這些訓練週期的速度。

每個訓練週期的處理時間都比預期更長 (數秒)，這表示 Chrome 已改回使用一般的 JS CPU 執行作業，而不是使用 GPU：

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

摘要

自 2017 年成立以來，網路 AI 已大幅成長。透過 WebGPU、WebGL 和 WebAssembly 等瀏覽器技術，機器學習模型的數學運算可在用戶端進一步加快。

截至 2023 年，TensorFlow.js 和 MediaPipe Web 超過 10 億下載模型和程式庫的下載次數。這是一種歷史性的里程碑，也是網頁開發人員和工程師如何轉向將 AI 技術導入新一代網頁應用程式，以打造出真正令人驚豔的解決方案的跡象。

如果能成功使用產品，就必須負起責任。在實際工作環境系統的這個用量等級中，需要在真正的瀏覽器環境中測試用戶端、瀏覽器式 AI 模型，同時具備可擴展性、可自動化的特性，以及在已知的標準化硬體配置中。

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

總結

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

如果您覺得這項功能很實用，歡迎透過 LinkedIn、X (舊稱 Twitter) 或任何社群網路 (使用 #WebAI 主題標記) 公開喊話。期待收到您的寶貴意見，讓我們日後能提供更多類似資訊。

在 GitHub 存放區上加上星號，即可接收日後發布的最新消息。

特別銘謝

感謝 Chrome 團隊的所有成員，也謝謝我們為本解決方案遇到的驅動程式和 WebGPU 問題進行偵錯。特別感謝 Jecelyn Yeen 和 Alexandra White 撰寫這篇網誌文章。感謝 Yuly Novikov、Andrey Kosyakov 及 Alex Rudenko，他在建立最終可用的解決方案時，全心投入。