Potenzia i test dei modelli di IA web: WebGPU, WebGL e Chrome headless

François Beaufort
François Beaufort

Ottime notizie: Hai creato un'applicazione di IA web interessante che esegue modelli di machine learning direttamente sul dispositivo di un utente. Viene eseguito interamente sul browser web lato client, senza fare affidamento sul cloud. Questa progettazione sul dispositivo migliora la privacy dell'utente, incrementa le prestazioni e riduce significativamente i costi.

Tuttavia, c'è un ostacolo. Il tuo modello TensorFlow.js può funzionare su entrambe le CPU (WebAssembly) e su GPU più potenti (tramite WebGL e WebGPU). La domanda è: come potete automatizzare in modo coerente i test del browser con l'hardware selezionato?

Il mantenimento della coerenza è fondamentale per confrontare le prestazioni dei modelli di machine learning nel tempo durante l'iterazione e il loro miglioramento, prima del deployment per gli utenti reali da utilizzare sul proprio dispositivo.

La configurazione di un ambiente di test coerente con GPU può essere più difficile del previsto. In questo post del blog, condivideremo i problemi che abbiamo affrontato e come li abbiamo risolti in modo che tu possa migliorare le prestazioni della tua applicazione.

Questo non è solo per gli sviluppatori di Web AI! Se ti occupi di giochi web o di immagini, questo post è prezioso anche per te.

Vantaggi dei nostri strumenti per l'automazione

Ecco cosa stiamo usando:

  • Ambiente: un blocco note Google Colab basato su Linux collegato a una GPU NVIDIA T4 o V100. Se preferisci, puoi utilizzare altre piattaforme cloud, come Google Cloud (Google Cloud).
  • Browser. Chrome supporta WebGPU, un potente successo di WebGL, che porta i progressi delle moderne API GPU sul web.
  • Automazione: Puppeteer è una libreria Node.js che ti consente di controllare i browser in modo programmatico con JavaScript. Con Puppeteer, possiamo automatizzare Chrome in modalità headless, il che significa che il browser viene eseguito senza un'interfaccia visibile su un server. Stiamo utilizzando la nuova modalità headless migliorata, non il modulo legacy.

Verifica l'ambiente

Il modo migliore per verificare se l'accelerazione hardware è attiva in Chrome è digitare chrome://gpu nella barra degli indirizzi. Puoi eseguire l'operazione equivalente con Puppeteer in modo programmatico con console.log o salvare il report completo in formato PDF per poterlo controllare manualmente:

/* 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();

Apri chrome://gpu. Dovresti vedere i seguenti risultati:

Stato della funzionalità delle immagini
OpenGL: Disabilitata
Vulkan: Disabilitata
WebGL: Solo software, accelerazione hardware non disponibile.
WebGL2: Solo software, accelerazione hardware non disponibile.
GPU Web: Disabilitata

Problemi rilevati.
WebGPU è stata disattivata tramite una lista bloccata o la riga di comando.

Non è un ottimo inizio. È abbastanza chiaro che il rilevamento dell'hardware non funzionava. WebGL, WebGL2 e WebGPU sono essenzialmente disabilitati o solo software. Non siamo soli a occuparci di questo problema: esistono numerose discussioni online di persone che vivono in una situazione simile, anche sui canali di assistenza ufficiali di Chrome (1), (2).

Attiva il supporto di WebGPU e WebGL

Per impostazione predefinita, Chrome headless disattiva GPU. Per abilitarla su Linux, applica tutti i seguenti flag all'avvio di Chrome Headless:

  • Il flag --no-sandbox disattiva la sandbox per la sicurezza di Chrome, che isola il processo del browser dal resto del sistema. Non è possibile eseguire Chrome come root senza questa sandbox.
  • Il flag --headless=new esegue Chrome con la nuova e migliorata modalità headless, senza UI visibile.
  • Il flag --use-angle=vulkan indica a Chrome di utilizzare il backend Vulkan per ANGLE, che traduce le chiamate OpenGL ES 2/3 in chiamate API Vulkan.
  • Il flag --enable-features=Vulkan abilita il backend Vulkan per la composizione e la rasterizzazione in Chrome.
  • Il flag --disable-vulkan-surface disabilita l'estensione VK_KHR_surface vulkan instance. Invece di utilizzare uno swapchain, Bit blit viene utilizzato per il risultato di rendering attuale sullo schermo.
  • Il flag --enable-unsafe-webgpu attiva l'API WebGPU sperimentale in Chrome su Linux e disattiva la lista bloccata degli adattatori.

Ora combiniamo tutte le modifiche apportate finora. Ecco lo script completo.

/* 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();

Esegui di nuovo lo script. Non vengono rilevati problemi con WebGPU e il valore passa da disattivato a solo software.

Stato della funzionalità delle immagini
OpenGL: Disabilitata
Vulkan: Disabilitata
WebGL: Solo software, accelerazione hardware non disponibile.
WebGL2: Solo software, accelerazione hardware non disponibile.
GPU Web: Solo software, accelerazione hardware non disponibile.

Tuttavia, l'accelerazione hardware non è ancora disponibile e la GPU NVIDIA T4 non viene rilevata.

Installa i driver GPU corretti

Abbiamo esaminato in modo più approfondito l'output di chrome://gpu, insieme ad alcuni esperti di GPU del team di Chrome. Sono stati rilevati problemi con i driver predefiniti installati nell'istanza Colab Linux, che causano problemi con Vulkan e impedivano a Chrome di rilevare la GPU NVIDIA T4 a livello GL_RENDERER, come mostrato nell'output seguente. Questo causa problemi con Chrome headless.

L'output predefinito non rileva la GPU NVIDIA T4.
Informazioni sul conducente
GL_RENDERER ANGLE (Google, Vulkan 1.3.0 (SwiftShader Device (Subzero)) (0x0000C0DE)), SwiftShader driver-5.0.0)

Di conseguenza, l'installazione dei driver corretti e compatibili consente di risolvere il problema.

Output aggiornato dopo l'installazione dei driver.
Informazioni sul conducente
GL_RENDERER ANGLE (NVIDIA Corporation, Tesla T4/PCIe/SSE2, OpenGL ES 3.2 NVIDIA 525.105.17)

Per installare i driver corretti, esegui questi comandi durante la configurazione. Le ultime due righe ti aiutano a registrare gli output dei driver NVIDIA rilevati insieme a vulkaninfo.

apt-get install -y vulkan-tools libnvidia-gl-525

// Verify the NVIDIA drivers detects along with vulkaninfo
nvidia-smi
vulkaninfo --summary

Ora esegui di nuovo lo script e ottieni il seguente risultato. 🎉

Stato della funzionalità delle immagini
OpenGL: Attivato
Vulkan: Attivato
WebGL: Accelerazione hardware, ma prestazioni ridotte.
WebGL2: Accelerazione hardware, ma prestazioni ridotte.
GPU Web: Accelerazione hardware, ma prestazioni ridotte.

Grazie all'utilizzo dei driver e dei flag corretti durante l'esecuzione di Chrome, ora supportiamo WebGPU e WebGL con la nuova modalità headless.

Dietro le quinte: l'indagine del nostro team

Dopo molte ricerche, non abbiamo trovato metodi di lavoro per l'ambiente che dovevamo eseguire in Google Colab, anche se erano presenti alcuni post di speranza che funzionavano in altri ambienti, il che era promettente. In definitiva, non siamo riusciti a replicare il loro successo nell'ambiente NVIDIA T4 di Colab, dato che avevamo due problemi principali:

  1. Alcune combinazioni di flag consentono il rilevamento della GPU, ma non ti consentono di utilizzare effettivamente la GPU.
  2. Esempi di soluzioni funzionanti di terze parti utilizzavano la vecchia versione headless di Chrome, che a un certo punto verrà deprecata in favore della nuova versione. Avevamo bisogno di una soluzione che funzionasse con il nuovo Chrome Headless per una migliore preparazione futura.

Abbiamo confermato l'utilizzo ridotto della GPU eseguendo una pagina web di esempio TensorFlow.js per il riconoscimento delle immagini, in cui abbiamo addestrato un modello per riconoscere campioni di abbigliamento (un po' come una sorta di "hello world" del machine learning).

Su una macchina normale, 50 cicli di addestramento (noti come epoche) dovrebbero essere eseguiti in meno di un secondo ciascuno. Chiamando Chrome Headless nello stato predefinito, potremmo registrare l'output della console JavaScript nella riga di comando lato server Node.js per vedere la velocità effettiva di questi cicli di addestramento.

Come previsto, ogni epoca di addestramento ha richiesto molto più tempo del previsto (diversi secondi), il che suggerisce che Chrome è tornato alla normale esecuzione della CPU JS invece di utilizzare la GPU:

Le epoche di addestramento si muovono a una cadenza più lenta.
Figura 1: acquisizione in tempo reale che mostra il tempo di esecuzione di ogni epoca di addestramento (secondi).

Dopo aver corretto i driver e utilizzato la giusta combinazione di flag per Chrome headless, rieseguire l'esempio di addestramento TensorFlow.js si traduce in epoche di addestramento molto più rapide.

C'è un aumento della velocità per le epoche...
Figura 2: acquisizione in tempo reale che mostra la velocità delle epoche.

Riepilogo

L'IA web è cresciuta in modo esponenziale da quando è stata creata nel 2017. Con le tecnologie browser come WebGPU, WebGL e WebAssembly, le operazioni matematiche di un modello di machine learning possono essere ulteriormente accelerate sul lato client.

Nel 2023 TensorFlow.js e MediaPipe Web hanno superato oltre 1 miliardo di download di modelli e librerie: un traguardo storico e un segno di come gli sviluppatori e gli ingegneri web stanno cambiando l'adozione dell'IA nelle loro app web di nuova generazione per realizzare soluzioni davvero incredibili.

Da un grande successo nell'utilizzo derivano grandi responsabilità. A questo livello di utilizzo nei sistemi di produzione, si pone la necessità di testare i modelli di IA lato client e basati su browser in un ambiente browser reale, garantendo al contempo la scalabilità, l'automazione e una configurazione hardware standardizzata nota.

Sfruttando la potenza combinata dei nuovi Chrome e Puppeteer Headless, puoi testare con sicurezza questi carichi di lavoro in un ambiente standardizzato e replicabile, garantendo risultati coerenti e affidabili.

Conclusione

Nella nostra documentazione è disponibile una guida passo passo, che ti consente di provare personalmente la configurazione completa.

Se lo hai trovato utile, ringrazialo su LinkedIn, X (in precedenza Twitter) o su qualsiasi altro social network che usi usando l'hashtag #WebAI. Mi farebbe piacere ricevere il tuo feedback per consentirci di scrivere altre cose simili in futuro.

Aggiungi una stella al repository GitHub per ricevere eventuali aggiornamenti futuri.

Ringraziamenti

Un enorme ringraziamento a tutti i membri del team di Chrome che ci hanno aiutato a eseguire il debug dei problemi relativi al driver e alle GPU WebGPU che abbiamo affrontato in questa soluzione, e un ringraziamento speciale a Jecelyn Yeen e Alexandra White per aver aiutato a scrivere questo post del blog. Grazie a Yuly Novikov, Andrey Kosyakov e Alex Rudenko, che sono stati determinanti nella creazione della soluzione finale funzionante.