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

François Beaufort
François Beaufort

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

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

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

設定 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,您應該會看到以下結果:

顯示卡功能狀態
OpenGL: 已停用
Vulkan: 已停用
WebGL: 僅限軟體,無法使用硬體加速。
WebGL2: 僅限軟體,無法使用硬體加速。
WebGPU: 已停用

偵測到問題。
已透過封鎖清單或指令列停用 WebGPU。

這不是個好的開始。很明顯,硬體偵測失敗。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 問題,且值已從「停用」變更為「僅限軟體」。

顯示卡功能狀態
OpenGL: 已停用
Vulkan: 已停用
WebGL: 僅限軟體,無法使用硬體加速。
WebGL2: 僅限軟體,無法使用硬體加速。
WebGPU: 僅限軟體,無法使用硬體加速。

不過,硬體加速功能仍無法使用,NVIDIA T4 GPU 未偵測到。

安裝正確的 GPU 驅動程式

我們與 Chrome 團隊的部分 GPU 專家一起,更仔細地調查 chrome://gpu 的輸出內容。我們發現在 Linux Colab 執行個體上安裝的預設驅動程式有問題,導致 Vulkan 發生問題,導致 Chrome 無法在 GL_RENDERER 層級偵測 NVIDIA T4 GPU,如以下輸出內容所示。這會導致 Headless 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 無頭版本,但這項版本將在某個時間點淘汰,改用新版。我們需要能與新版無頭 Chrome 搭配使用的解決方案,以便因應未來需求。

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

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

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

訓練週期會以較慢的節奏進行。
圖 1:即時擷取畫面,顯示每個訓練迴圈的執行時間 (秒)。

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

訓練週期的速度會加快。
圖 2:即時擷取畫面,顯示各個 epoch 的加速效果。

摘要

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

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

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

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

總結

我們的文件中提供逐步指南,您可以自行嘗試完整的設定程序。

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

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

特別銘謝

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