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

Jason Mayes

François Beaufort

Ottime notizie! Hai creato una fantastica applicazione di IA web 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. Questo design sul dispositivo migliora la privacy degli utenti, migliora le prestazioni e riduce notevolmente i costi.

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

Mantenere la coerenza è fondamentale per confrontare il rendimento del modello di machine learning nel tempo man mano che lo esegui l'iterazione e lo migliori, prima di eseguirlo per gli utenti reali sul loro dispositivo.

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

Non è solo per gli sviluppatori di IA web. Se ti occupi di giochi o grafica web, questo post è utile anche per te.

Che cosa contiene la nostra cassetta degli attrezzi per l'automazione

Ecco cosa utilizziamo:

• Ambiente: un notebook Google Colab basato su Linux collegato a una GPU NVIDIA T4 o V100. Se preferisci, puoi utilizzare altre piattaforme cloud, come Google Cloud.
• Browser: Chrome supporta WebGPU, un potente successore di WebGL che porta sul web i progressi delle API GPU moderne.
• Automazione: Puppeteer è una libreria Node.js che consente di controllare i browser in modo programmatico con JavaScript. Con Puppeteer, possiamo eseguire l'automazione di Chrome in modalità headless, il che significa che il browser viene eseguito su un server senza un'interfaccia visibile. Utilizziamo 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'equivalente con Puppeteer in modo programmatico con console.log o salvare il report completo 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 e dovresti visualizzare i seguenti risultati:

Non è un buon inizio. È abbastanza chiaro che il rilevamento dell'hardware non è riuscito. WebGL, WebGL2 e WebGPU sono sostanzialmente disattivati o solo software. Non siamo soli a riscontrare questo problema: online sono presenti numerose discussioni di persone in una situazione simile, anche sui canali di assistenza ufficiali di Chrome (1), (2).

Attivare il supporto di WebGPU e WebGL

Per impostazione predefinita, Headless Chrome disattiva la GPU. Per attivarlo su Linux, applica tutti i flag riportati di seguito quando avvii 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 utente root senza questa sandbox non è supportata.
• Il flag --headless=new esegue Chrome con la nuova e migliorata modalità headless, senza alcuna interfaccia utente 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 all'API Vulkan.
• Il flag --enable-features=Vulkan attiva il backend grafico Vulkan per il compositing e la rasterizzazione in Chrome.
• Il flag --disable-vulkan-surface disattiva l'estensione dell'istanza VK_KHR_surface vulkan. Anziché utilizzare una swapchain, viene utilizzato il blit di bit per visualizzare il risultato del 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 a WebGPU e il valore passa da disabilitato a solo software.

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

Installa i driver GPU corretti

Abbiamo esaminato più da vicino l'output di chrome://gpu con alcuni esperti di GPU del team di Chrome. Abbiamo riscontrato problemi con i driver predefiniti installati sull'istanza Linux Colab, che causano problemi con Vulkan, impedendo a Chrome di rilevare la GPU NVIDIA T4 a livello GL_RENDERER, come mostrato nell'output seguente. Ciò causa problemi con Headless Chrome.

L'installazione dei driver corretti e compatibili risolve quindi il problema.

Per installare i driver corretti, esegui i seguenti comandi durante la configurazione. Le due righe finali ti aiutano a registrare le uscite di ciò che viene rilevato dai driver NVIDIA 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. 🎉

Utilizzando i driver e i flag corretti durante l'esecuzione di Chrome, ora abbiamo il supporto di WebGPU e WebGL utilizzando la nuova modalità headless.

Dietro le quinte: l'indagine del nostro team

Dopo molte ricerche, non abbiamo trovato metodi funzionanti per l'ambiente che dovevamo eseguire in Google Colab, anche se alcuni post incoraggianti funzionavano in altri ambienti, il che era promettente. Alla fine, non siamo riusciti a replicare il loro successo nell'ambiente Colab NVIDIA T4, in quanto abbiamo riscontrato due problemi principali:

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

Abbiamo confermato il sottoutilizzo della GPU eseguendo un'pagina web di esempio di TensorFlow.js per il riconoscimento delle immagini, con la quale abbiamo addestrato un modello per riconoscere campioni di abbigliamento (un po' come un "ciao mondo" del machine learning).

Su una macchina normale, 50 cicli di addestramento (chiamati epoche) devono essere eseguiti in meno di 1 secondo ciascuno. Chiamando Headless Chrome nel suo stato predefinito, abbiamo potuto registrare l'output della console JavaScript nella riga di comando lato server di Node.js per vedere quanto erano effettivamente rapidi 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 all'esecuzione della CPU JS normale anziché utilizzare la GPU:

Dopo aver corretto i driver e aver utilizzato la combinazione corretta di flag per Chrome headless, la riesecuzione dell'esempio di addestramento di TensorFlow.js genera epoche di addestramento molto più rapide.

Riepilogo

L'IA web è cresciuta in modo esponenziale dalla sua creazione nel 2017. Con le tecnologie del browser come WebGPU, WebGL e WebAssembly, le operazioni matematiche di un modello di machine learning possono essere ulteriormente accelerate lato client.

A partire dal 2023, TensorFlow.js e MediaPipe Web hanno superato 1 miliardo di download di modelli e librerie, un traguardo storico e un segno di come gli sviluppatori e gli ingegneri web stiano adottando l'IA nelle loro app web di nuova generazione per creare soluzioni davvero incredibili.

Un grande successo nell'utilizzo comporta una grande responsabilità. A questo livello di utilizzo nei sistemi di produzione, è necessario testare i modelli di IA lato client e basati su browser in un vero ambiente browser, ma che siano anche scalabili, automatizzabili e all'interno di una configurazione hardware standardizzata nota.

Sfruttando la potenza combinata del nuovo Chrome Headless e di Puppeteer, 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 la configurazione completa.

Se hai trovato utile questa risorsa, fai un fischio su LinkedIn, X (in precedenza Twitter) o su qualsiasi social network tu utilizzi utilizzando l'hashtag #WebAI. Sarebbe fantastico ricevere il tuo feedback, in modo da sapere cosa scrivere in futuro.

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

Ringraziamenti

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