Supercharge Web AI-modeltesten: WebGPU, WebGL en Headless Chrome

Jason Mayes

François Beaufort

Geweldig nieuws! Je hebt een coole Web AI-applicatie gebouwd die machine learning-modellen rechtstreeks op het apparaat van een gebruiker uitvoert. Het draait volledig op de webbrowser aan de clientzijde, zonder afhankelijk te zijn van de cloud. Dit ontwerp op het apparaat verbetert de privacy van gebruikers, verbetert de prestaties en verlaagt de kosten aanzienlijk.

Er is echter een hindernis. Uw TensorFlow.js -model kan werken op zowel CPU's (WebAssembly) als krachtigere GPU's (via WebGL en WebGPU ). De vraag is: hoe kun je het testen van browsers consistent automatiseren met de geselecteerde hardware?

Het behouden van consistentie is van cruciaal belang voor het vergelijken van de prestaties van machine learning-modellen in de loop van de tijd terwijl u deze herhaalt en verbetert, voordat ze worden geïmplementeerd voor echte gebruikers voor gebruik op hun apparaat.

Het opzetten van een consistente testomgeving met GPU's kan lastiger zijn dan verwacht. In deze blogpost delen we de problemen waarmee we te maken kregen en hoe we deze hebben opgelost, zodat u de prestaties van uw applicatie kunt verbeteren.

Dit is niet alleen voor Web AI-ontwikkelaars! Als je aan webgaming of graphics werkt, is dit bericht ook waardevol voor jou.

Wat zit er in onze automatiseringstoolbox

Dit is wat we gebruiken:

• Omgeving : een op Linux gebaseerde Google Colab- notebook aangesloten op een NVIDIA T4 of V100 GPU. Indien gewenst kunt u andere cloudplatforms gebruiken, zoals Google Cloud (GCP).
• Browser : Chrome ondersteunt WebGPU , een krachtige opvolger van WebGL , die de vooruitgang van moderne GPU-API's naar internet brengt.
• Automatisering : Puppeteer is een Node.js-bibliotheek waarmee u browsers programmatisch kunt besturen met JavaScript. Met Puppeteer kunnen we Chrome automatiseren in headless-modus, wat betekent dat de browser zonder zichtbare interface op een server draait. We gebruiken de verbeterde nieuwe headless-modus , niet het oude formulier.

Controleer de omgeving

De beste manier om te controleren of hardwareversnelling in Chrome is ingeschakeld, is door chrome://gpu in de adresbalk te typen. U kunt het equivalent programmatisch uitvoeren met Puppeteer met console.log of het volledige rapport opslaan als PDF om het handmatig te controleren:

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

Open chrome://gpu en je zou de volgende resultaten moeten krijgen:

Geen goed begin. Het is vrij duidelijk dat de hardwaredetectie faalde. WebGL, WebGL2 en WebGPU zijn in wezen uitgeschakeld of alleen software . We zijn niet de enige met dit probleem: er zijn online talloze discussies van mensen in een vergelijkbare situatie, ook op de officiële Chrome-ondersteuningskanalen ( 1 ), ( 2 ).

Schakel WebGPU- en WebGL-ondersteuning in

Standaard schakelt Headless Chrome GPU uit . Om het op Linux in te schakelen, past u alle volgende vlaggen toe bij het starten van Headless Chrome:

• --no-sandbox flag schakelt de beveiligingssandbox van Chrome uit, die het browserproces isoleert van de rest van het systeem. Het uitvoeren van Chrome als root zonder deze sandbox wordt niet ondersteund.
• --headless=new vlag voert Chrome uit met de nieuwe en verbeterde headless-modus , zonder enige zichtbare gebruikersinterface.
• --use-angle=vulkan flag vertelt Chrome om de Vulkan backend voor ANGLE te gebruiken, wat OpenGL ES 2/3-aanroepen vertaalt naar Vulkan API-aanroepen.
• --enable-features=Vulkan vlag schakelt Vulkan grafische backend in voor compositie en rasterisatie in Chrome.
• --disable-vulkan-surface vlag schakelt de VK_KHR_surface vulkan instantie-extensie uit. In plaats van een swapchain te gebruiken, wordt Bit blit gebruikt voor het huidige renderresultaat op het scherm.
• --enable-unsafe-webgpu flag schakelt de experimentele WebGPU API in Chrome op Linux in en schakelt de adapterblokkeringslijst uit.

Nu combineren we alle wijzigingen die we tot nu toe hebben aangebracht. Hier is het volledige script.

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

Voer het script opnieuw uit. Er worden geen WebGPU-problemen gedetecteerd en de waarde verandert van uitgeschakeld naar alleen software.

Hardwareversnelling is echter nog steeds niet beschikbaar en de NVIDIA T4 GPU wordt niet gedetecteerd.

Installeer de juiste GPU-stuurprogramma's

We hebben de uitvoer van chrome://gpu nader onderzocht, met enkele GPU-experts in het Chrome-team. We hebben problemen gevonden met de standaardstuurprogramma's die op de Linux Colab-instantie zijn geïnstalleerd , waardoor problemen met Vulkan ontstonden, waardoor Chrome de NVIDIA T4 GPU op GL_RENDERER -niveau niet kon detecteren, zoals weergegeven in de volgende uitvoer. Dit veroorzaakt problemen met Headless Chrome.

Het installeren van de juiste stuurprogramma's die compatibel waren, lost daarom het probleem op.

Om de juiste stuurprogramma's te installeren, voert u tijdens de installatie de volgende opdrachten uit. De laatste twee regels helpen je bij het loggen van de uitvoer van wat NVIDIA-stuurprogramma's samen met vulkaninfo detecteren.

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

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

Voer het script nu opnieuw uit en we krijgen het volgende resultaat. 🎉

Door de juiste stuurprogramma's en vlaggen te gebruiken bij het uitvoeren van Chrome, hebben we nu WebGPU- en WebGL-ondersteuning met behulp van de glanzende, nieuwe headless-modus.

Achter de schermen: het onderzoek van ons team

Na veel onderzoek vonden we geen werkmethoden voor de omgeving die we moesten uitvoeren in Google Colab, hoewel er enkele hoopvolle berichten waren die in andere omgevingen werkten, wat veelbelovend was. Uiteindelijk konden we hun succes niet repliceren in de Colab NVIDIA T4-omgeving, omdat we twee belangrijke problemen hadden:

1. Sommige combinaties van vlaggen maken detectie van de GPU mogelijk, maar zorgen ervoor dat u de GPU niet daadwerkelijk kunt gebruiken.
2. Voorbeelden van werkende oplossingen van derden gebruikten de oude headless-versie van Chrome, die op een gegeven moment verouderd zal zijn ten gunste van de nieuwe versie . We hadden een oplossing nodig die werkte met het nieuwe Headless Chrome om beter toekomstbestendig te zijn.

We hebben de onderbenutting van de GPU bevestigd door een voorbeeld van een TensorFlow.js-webpagina voor beeldherkenning uit te voeren, waarbij we een model hebben getraind om kledingstalen te herkennen (een beetje zoals een "hallo wereld" van machinaal leren).

Op een gewone machine moeten 50 trainingscycli (ook wel tijdperken genoemd) elk in minder dan 1 seconde worden uitgevoerd. Door Headless Chrome in de standaardstatus aan te roepen, konden we de uitvoer van de JavaScript-console loggen op de Node.js-opdrachtregel aan de serverzijde om te zien hoe snel deze trainingscycli feitelijk duurden.

Zoals verwacht duurde elk trainingstijdvak veel langer dan verwacht (enkele seconden), wat suggereert dat Chrome is teruggevallen naar de gewone JS CPU-uitvoering in plaats van de GPU te gebruiken:

Na het repareren van de stuurprogramma's en het gebruiken van de juiste combinatie van vlaggen voor Headless Chrome, resulteert het opnieuw uitvoeren van het TensorFlow.js-trainingsvoorbeeld in veel snellere trainingsperioden.

Samenvatting

Web AI is exponentieel gegroeid sinds de oprichting ervan in 2017. Met browsertechnologieën zoals WebGPU, WebGL en WebAssembly kunnen de wiskundige bewerkingen van een machine learning-model aan de clientzijde verder worden versneld.

Vanaf 2023 hebben TensorFlow.js en MediaPipe Web meer dan 1 miljard downloads van modellen en bibliotheken gekruist – een historische mijlpaal en een teken van hoe webontwikkelaars en -ingenieurs AI in hun volgende generatie webapps gaan omarmen om werkelijk ongelooflijke oplossingen te maken .

Met groot succes in het gebruik komt een grote verantwoordelijkheid. Op dit gebruiksniveau in productiesystemen ontstaat de behoefte aan het testen van browsergebaseerde AI-modellen aan de clientzijde in een echte browseromgeving, terwijl ze ook schaalbaar en automatiseerbaar zijn en binnen een bekende gestandaardiseerde hardware-installatie.

Door de gecombineerde kracht van de nieuwe Headless Chrome en Puppeteer te benutten, kunt u dergelijke werklasten met vertrouwen testen in een gestandaardiseerde en repliceerbare omgeving, waardoor consistente en betrouwbare resultaten worden gegarandeerd.

Sluit af

In onze documentatie vindt u een stapsgewijze handleiding , zodat u de volledige installatie zelf kunt uitproberen.

Als je dit nuttig vond, stuur dan een bericht naar LinkedIn , X (voorheen Twitter) of welk sociaal netwerk je ook gebruikt met de hashtag #WebAI . Het zou geweldig zijn om uw feedback te horen, zodat we weten dat we in de toekomst meer van dit soort dingen kunnen schrijven.

Voeg een ster toe aan de Github-repository om toekomstige updates te ontvangen.

Dankbetuigingen

Hartelijk dank aan iedereen in het Chrome-team die heeft geholpen bij het oplossen van de driver- en WebGPU-problemen waarmee we bij deze oplossing te maken kregen, met een speciale dank aan Jecelyn Yeen en Alexandra White voor hun hulp bij het schrijven van deze blogpost. Met dank aan Yuly Novikov, Andrey Kosyakov en Alex Rudenko die een belangrijke rol hebben gespeeld bij het creëren van de uiteindelijke, werkende oplossing.