Optimiser les tests de modèles d'IA Web: GPU Web, WebGL et Chrome sans interface graphique

Jason Mayes

François Beaufort

Bonne nouvelle ! Vous avez créé une application Web d'IA qui exécute des modèles de machine learning directement sur l'appareil d'un utilisateur. Elle s'exécute entièrement dans le navigateur Web côté client, sans s'appuyer sur le cloud. Cette conception sur l'appareil améliore la confidentialité des utilisateurs, améliore les performances et réduit considérablement les coûts.

Il y a toutefois un obstacle. Votre modèle TensorFlow.js peut fonctionner à la fois sur des CPU (WebAssembly) et sur des GPU plus puissants (via WebGL et WebGPU). La question est la suivante : comment pouvez-vous automatiser de manière cohérente les tests de navigateur avec le matériel sélectionné ?

Il est essentiel de maintenir la cohérence pour comparer les performances des modèles de machine learning au fil du temps à mesure que vous les itérez et les améliorez, avant de les déployer pour que les utilisateurs réels puissent les utiliser sur leur appareil.

La configuration d'un environnement de test cohérent avec des GPU peut être plus difficile que prévu. Dans cet article de blog, nous allons partager les problèmes auxquels nous avons été confrontés et la façon dont nous les avons résolus, afin que vous puissiez améliorer les performances de votre application.

Ce n'est pas seulement pour les développeurs d'IA Web ! Si vous travaillez sur les jeux ou les graphismes Web, cet article vous sera également utile.

Contenu de notre boîte à outils d'automatisation

Voici ce que nous utilisons:

• Environnement: notebook Google Colab basé sur Linux connecté à un GPU NVIDIA T4 ou V100. Si vous préférez, vous pouvez utiliser d'autres plates-formes cloud, telles que Google Cloud (GCP).
• Navigateur: Chrome est compatible avec WebGPU, un puissant successeur de WebGL qui apporte les avancées des API GPU modernes sur le Web.
• Automatisation: Puppeteer est une bibliothèque Node.js qui vous permet de contrôler les navigateurs de manière programmatique avec JavaScript. Avec Puppeteer, nous pouvons automatiser Chrome en mode headless, ce qui signifie que le navigateur s'exécute sans interface visible, sur un serveur. Nous utilisons le nouveau mode headless amélioré, et non l'ancien.

Vérifier l'environnement

Le meilleur moyen de vérifier si l'accélération matérielle est activée dans Chrome consiste à saisir chrome://gpu dans la barre d'adresse. Vous pouvez effectuer l'équivalent avec Puppeteer de manière programmatique avec console.log ou enregistrer le rapport complet au format PDF pour le vérifier manuellement:

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

Ouvrez chrome://gpu. Vous devriez obtenir les résultats suivants:

Ce n'est pas un bon début. Il est assez clair que la détection matérielle a échoué. WebGL, WebGL2 et WebGPU sont essentiellement désactivés ou logiciels uniquement. Nous ne sommes pas les seuls à rencontrer ce problème. De nombreuses discussions en ligne font état de situations similaires, y compris sur les canaux d'assistance Chrome officiels (1), (2).

Activer la compatibilité avec WebGPU et WebGL

Par défaut, Chrome headless désactive le GPU. Pour l'activer sous Linux, appliquez tous les indicateurs suivants lorsque vous lancez Chrome sans tête:

• L'indicateur --no-sandbox désactive le bac à sable de sécurité de Chrome, qui isole le processus du navigateur du reste du système. L'exécution de Chrome en tant que root sans ce bac à sable n'est pas prise en charge.
• L'indicateur --headless=new exécute Chrome avec le nouveau et amélioré mode headless, sans interface utilisateur visible.
• L'indicateur --use-angle=vulkan indique à Chrome d'utiliser le backend Vulkan pour ANGLE, qui traduit les appels OpenGL ES 2/3 en appels d'API Vulkan.
• L'indicateur --enable-features=Vulkan active le backend graphique Vulkan pour le compositing et la rastérisation dans Chrome.
• L'indicateur --disable-vulkan-surface désactive l'extension d'instance Vulkan VK_KHR_surface. Au lieu d'utiliser un swapchain, le blit de bits est utilisé pour le résultat de rendu actuel à l'écran.
• Le flag --enable-unsafe-webgpu active l'API WebGPU expérimentale dans Chrome sur Linux et désactive la liste de blocage des adaptateurs.

Nous allons maintenant combiner toutes les modifications que nous avons apportées jusqu'à présent. Voici le script complet :

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

Exécutez à nouveau le script. Aucun problème WebGPU n'est détecté, et la valeur passe de "désactivé" à "logiciel uniquement".

Toutefois, l'accélération matérielle n'est toujours pas disponible, et le GPU NVIDIA T4 n'est pas détecté.

Installer les pilotes de GPU appropriés

Nous avons examiné plus en détail la sortie de chrome://gpu, avec certains experts GPU de l'équipe Chrome. Nous avons détecté des problèmes avec les pilotes par défaut installés sur l'instance Linux Colab, qui ont entraîné des problèmes avec Vulkan, ce qui a empêché Chrome de détecter le GPU NVIDIA T4 au niveau GL_RENDERER, comme indiqué dans la sortie suivante. Cela entraîne des problèmes avec Chrome headless.

L'installation des pilotes appropriés et compatibles résout donc le problème.

Pour installer les pilotes appropriés, exécutez les commandes suivantes lors de la configuration. Les deux dernières lignes vous aident à consigner les sorties de ce que les pilotes NVIDIA détectent avec vulkaninfo.

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

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

Exécutez à nouveau le script. Vous obtenez le résultat suivant. 🎉

En utilisant les pilotes et les indicateurs appropriés lorsque vous exécutez Chrome, vous bénéficiez désormais de la prise en charge de WebGPU et de WebGL à l'aide du nouveau mode headless.

En coulisses: l'enquête de notre équipe

Après de nombreuses recherches, nous n'avons pas trouvé de méthodes fonctionnelles pour l'environnement que nous devions exécuter dans Google Colab, bien que certains posts encourageants aient fonctionné dans d'autres environnements, ce qui était prometteur. Finalement, nous n'avons pas pu reproduire leur succès dans l'environnement Colab NVIDIA T4, car nous rencontrions deux problèmes clés:

1. Certaines combinaisons d'indicateurs permettent de détecter le GPU, mais ne vous permettent pas de l'utiliser.
2. Les exemples de solutions fonctionnelles de tiers utilisaient l'ancienne version headless de Chrome, qui sera à terme abandonnée au profit de la nouvelle version. Nous avions besoin d'une solution compatible avec le nouveau Chrome sans interface utilisateur pour mieux nous préparer à l'avenir.

Nous avons confirmé le sous-utilisation du GPU en exécutant un exemple de page Web TensorFlow.js pour la reconnaissance d'images, dans lequel nous avons entraîné un modèle pour reconnaître des échantillons de vêtements (un peu comme un "Bonjour le monde" du machine learning).

Sur une machine standard, 50 cycles d'entraînement (appelés "épopées") doivent s'exécuter en moins d'une seconde chacun. En appelant Chrome sans tête dans son état par défaut, nous avons pu consigner la sortie de la console JavaScript dans la ligne de commande côté serveur Node.js pour voir à quelle vitesse ces cycles d'entraînement étaient réellement effectués.

Comme prévu, chaque époque d'entraînement a pris beaucoup plus de temps que prévu (plusieurs secondes), ce qui suggère que Chrome est revenu à l'exécution du processeur JS au lieu d'utiliser le GPU:

Après avoir corrigé les pilotes et utilisé la bonne combinaison d'indicateurs pour Chrome sans tête, relancer l'exemple d'entraînement TensorFlow.js génère des époques d'entraînement beaucoup plus rapides.

Résumé

L'IA Web a connu une croissance exponentielle depuis sa création en 2017. Grâce aux technologies de navigateur telles que WebGPU, WebGL et WebAssembly, les opérations mathématiques d'un modèle de machine learning peuvent être encore accélérées côté client.

En 2023, TensorFlow.js et MediaPipe Web ont dépassé le milliard de téléchargements de modèles et de bibliothèques. Il s'agit d'un jalon historique et d'un signe de l'adoption croissante de l'IA dans les applications Web de nouvelle génération pour créer des solutions vraiment incroyables.

Un grand succès s'accompagne de grandes responsabilités. À ce niveau d'utilisation dans les systèmes de production, il est nécessaire de tester les modèles d'IA côté client basés sur le navigateur dans un véritable environnement de navigateur, tout en étant évolutifs, automatisables et dans une configuration matérielle standardisée connue.

En exploitant la puissance combinée de la nouvelle version de Chrome sans interface utilisateur et de Puppeteer, vous pouvez tester en toute confiance ces charges de travail dans un environnement standardisé et reproductible, ce qui garantit des résultats cohérents et fiables.

Conclusion

Un guide détaillé est disponible dans notre documentation pour que vous puissiez essayer la configuration complète par vous-même.

Si vous avez trouvé cet article utile, n'hésitez pas à nous en parler sur LinkedIn, X (anciennement Twitter) ou sur le réseau social de votre choix en utilisant le hashtag #WebAI. Nous aimerions connaître votre avis pour savoir si nous devons rédiger d'autres articles de ce type à l'avenir.

Ajoutez une étoile au dépôt GitHub pour recevoir les futures mises à jour.

Remerciements

Un grand merci à tous les membres de l'équipe Chrome qui ont contribué à déboguer les problèmes de pilote et de WebGPU que nous avons rencontrés dans cette solution, et un merci tout particulier à Jecelyn Yeen et à Alexandra White pour leur aide à la rédaction de cet article de blog. Merci à Yuly Novikov, Andrey Kosyakov et Alex Rudenko, qui ont été essentiels à la création de la solution finale.