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

Ottime notizie: Hai creato un'applicazione IA web interessante che esegue modelli di machine learning direttamente sul dispositivo dell'utente. Viene eseguita interamente sul browser web lato client, senza fare affidamento sul cloud. Questo design on-device migliora la privacy dell'utente, incrementa le prestazioni e riduce in modo significativo i costi.

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

Mantenere la coerenza è fondamentale per confrontare le prestazioni dei modelli di machine learning nel tempo durante l'iterazione e il miglioramento, prima del deployment che gli utenti reali potranno utilizzare sul proprio dispositivo.

Configurare un ambiente di test coerente con le GPU può essere più difficile del previsto. In questo post del blog, condivideremo i problemi riscontrati e come li abbiamo risolti per consentirti di migliorare le prestazioni della tua applicazione.

Questo non è solo per gli sviluppatori di IA web. Se ti occupi di giochi sul web o di immagini grafiche, questo post è importante anche per te.

Cosa c'è nel nostro set di strumenti di automazione

Ecco cosa stiamo utilizzando:

• Ambiente: un notebook Google Colab basato su Linux connesso a una GPU NVIDIA T4 o V100. Puoi utilizzare altre piattaforme cloud, come Google Cloud (GCP), se preferisci.
• Browser. Chrome supporta WebGPU, un potente successo di WebGL, che porta sul web i progressi delle moderne API GPU.
• Automazione: Puppeteer è una libreria Node.js che 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 in modo programmatico eseguire l'equivalente con Puppeteer con console.log o salvare l'intero report come PDF per controllarlo 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 ottenere i seguenti risultati:

Stato della funzionalità della grafica
OpenGL:	Disabilitata
Vulkan:	Disabilitata
WebGL:	Solo software, accelerazione hardware non disponibile.
WebGL2:	Solo software, accelerazione hardware non disponibile.
GPU Web:	Disabilitata
Problemi rilevati. WebGPU è stato disattivato tramite lista bloccata o riga di comando.

Non un ottimo inizio. È abbastanza chiaro che il rilevamento dell'hardware non è riuscito. WebGL, WebGL2 e WebGPU sono essenzialmente disattivati o sono solo software. Non siamo soli in questo problema: ci sono numerose discussioni online di persone 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 senza testa disattiva le 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. L'esecuzione di Chrome come root senza questa sandbox non è supportata.
• Il flag --headless=new esegue Chrome con la nuova modalità headless migliorata, senza UI visibile.
• Il flag --use-angle=vulkan indica a Chrome di usare 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 grafico Vulkan per la composizione e la rasterizzazione in Chrome.
• Il flag --disable-vulkan-surface disabilita l'estensione dell'istanza vulkan VK_KHR_surface. Anziché utilizzare una swapchain, viene utilizzato Bit blit per il risultato di rendering 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 relativi alla WebGPU e il valore cambia da disabilitato a solo software.

Stato della funzionalità della grafica
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 ancora non è 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, con alcuni esperti di GPU del team di Chrome. Sono stati rilevati problemi relativi ai driver predefiniti installati sull'istanza Colab di Linux, che causano problemi con Vulkan e hanno portato Chrome a non riuscire a rilevare la GPU NVIDIA T4 a livello GL_RENDERER, come mostrato nel seguente output. 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 (Dispositivo SwiftShader (Subzero) (0x0000C0DE)), driver SwiftShader 5.0.0)

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 i comandi seguenti durante la configurazione. Le ultime due righe consentono di registrare gli output di ciò che i driver NVIDIA rilevano 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 otteniamo il seguente risultato. 🎉

Stato della funzionalità della grafica
OpenGL:	Abilitata
Vulkan:	Abilitata
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 possiamo contare sul supporto di WebGPU e WebGL che utilizzano la nuova e lucida 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 c'erano alcuni post sperati che funzionavano in altri ambienti, il che era promettente. Alla fine, non siamo riusciti a replicare il loro successo nell'ambiente Colab NVIDIA T4, a causa di due problemi chiave:

1. Alcune combinazioni di flag consentono il rilevamento della GPU, ma non ti permettono di utilizzare effettivamente la GPU.
2. Esempi di soluzioni funzionanti di terze parti utilizzano la vecchia versione headless di Chrome, che a un certo punto verrà ritirata a favore della nuova versione. Avevamo bisogno di una soluzione che funzionasse con il nuovo Chrome senza testa per essere più a prova di futuro.

Abbiamo verificato l'utilizzo insufficiente della GPU eseguendo una pagina web TensorFlow.js di esempio per il riconoscimento delle immagini, in cui abbiamo addestrato un modello a riconoscere campioni di abbigliamento (una specie di "Hello World" del machine learning).

Su una macchina normale, 50 cicli di addestramento (epoche) dovrebbero essere eseguiti in meno di un secondo ciascuno. Chiamando Chrome headless nello stato predefinito, abbiamo potuto registrare l'output della console JavaScript nella riga di comando lato server di Node.js per verificare 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 precedente anziché utilizzare la GPU:

Dopo aver corretto i driver e utilizzato la giusta combinazione di flag per Chrome senza testa, eseguire nuovamente l'esempio di addestramento TensorFlow.js porta a epoche di addestramento molto più rapide.

Riepilogo

L'IA per il web è cresciuta in modo esponenziale dalla sua creazione nel 2017. Con 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 del modo in cui gli sviluppatori web e gli ingegneri stanno passando all'IA nelle loro app web di nuova generazione per realizzare soluzioni davvero incredibili.

Un grande successo nell'utilizzo comporta una grande responsabilità. A questo livello di utilizzo nei sistemi di produzione, nasce la necessità di testare modelli di IA lato client basati su browser in un ambiente browser reale, oltre a essere scalabili, automatizzabili e all'interno di una configurazione hardware standardizzata nota.

Sfruttando la potenza combinata dei nuovi Chrome e Puppeteer senza testa, puoi testare con sicurezza tali 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 permetterà di provare in autonomia la configurazione completa.

Se l'hai trovato utile, scrivi un ringraziamento su LinkedIn, X (in precedenza Twitter) o su qualsiasi social network che utilizzi, usando l'hashtag #WebAI. Mi piacerebbe ricevere il tuo feedback, per consentirci di scrivere altri contenuti come questo in futuro.

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

Ringraziamenti

Un sentito ringraziamento a tutti i membri del team di Chrome che hanno contribuito a eseguire il debug dei problemi relativi a driver e WebGPU riscontrati in questa soluzione, con un ringraziamento speciale a Jecelyn Yeen e Alexandra White per l'aiuto a Wordsmith per questo post del blog. Grazie a Yuly Novikov, Andrey Kosyakov e Alex Rudenko che sono stati fondamentali nella creazione della soluzione definitiva e funzionante.