JAX

Publié le 26 octobre 2023. Google a mis à disposition du grand public Metrax, une nouvelle bibliothèque open source pour JAX, destinée à simplifier et à optimiser l’évaluation des modèles d’apprentissage automatique, un défi crucial pour les équipes travaillant à grande échelle.

• Metrax offre des implémentations standardisées et performantes de métriques courantes pour la classification, la régression, le traitement du langage naturel (TLN), la vision et l’audio.
• La bibliothèque prend en charge les environnements de formation distribués et à grande échelle, et utilise des fonctionnalités avancées de JAX pour améliorer les performances.
• Metrax vise à garantir la fiabilité des mesures et à standardiser les pratiques d’évaluation, un besoin exprimé par les développeurs migrant de TensorFlow vers JAX.

JAX, une bibliothèque Python open source pour le calcul numérique haute performance et l’apprentissage automatique, manquait jusqu’à présent d’un ensemble complet d’outils standardisés pour l’évaluation des modèles. De nombreuses équipes, lors de leur transition depuis TensorFlow, se sont retrouvées contraintes de développer leurs propres versions de métriques essentielles telles que la précision, le F1-score ou l’erreur quadratique moyenne (RMS). Metrax vient combler cette lacune, en fournissant une solution robuste et optimisée.

La bibliothèque propose des métriques prédéfinies pour divers types de modèles, allant de la classification à la recommandation, en passant par la vision artificielle et l’analyse audio. Pour les applications de vision, Metrax inclut des indicateurs tels que l’intersection sur union (IoU), le rapport signal sur bruit (SNR) et l’indice de similarité structurelle (SSIM). Dans le domaine du TLN, elle prend en charge des métriques comme la perplexité, BLEU et ROUGE.

Google souligne l’importance d’une implémentation rigoureuse et conforme aux meilleures pratiques pour ces mesures. Metrax exploite les capacités avancées de JAX, notamment les fonctions vmap et jit, pour maximiser les performances. Par exemple, la mise en œuvre des métriques « at K » permet de calculer plusieurs valeurs de K en parallèle, accélérant ainsi le processus d’évaluation et offrant une analyse plus complète.

« Vous pouvez utiliser PrecisionAtK pour déterminer la précision de votre modèle pour plusieurs valeurs de K (disons, à K=1, K=8 et K=20), le tout en un seul passage dans votre modèle, plutôt que d’avoir besoin d’appeler PrecisionAtK plusieurs fois avec chacun de ces arguments. »

Google

Un ingénieur DevOps, qui écrit sur Substack sous le pseudonyme Neural Foundry, a salué cette initiative :

« Le fait que Metrax prenne en charge le calcul de plusieurs valeurs K en un seul passage constitue une énorme victoire pour les systèmes de classement. J’ai réécrit les utilitaires de métriques à chaque fois que je change de projet et ce type de standardisation est attendu depuis longtemps. L’API a également l’air propre. Je suis curieux de savoir s’ils l’ont comparé à des implémentations personnalisées pour des cas d’utilisation spécifiques tels que des pipelines de recommandations à grande échelle. »

Neural Foundry, ingénieur DevOps

Voici un exemple de code Python illustrant comment calculer les mesures de précision à l’aide de Metrax :

import metrax

# Calcul direct de l'état de la métrique.
metric_state = metrax.Precision.from_model_output(
    predictions=predictions,
    labels=labels,
    threshold=0.5
)

# Le résultat est ensuite facilement accessible en appelant compute().
result = metric_state.compute()
result

En complément de la bibliothèque, Google a également publié un cahier d’exercices contenant des exemples détaillés, notamment sur la mise à l’échelle multi-appareils et l’intégration avec Flax, une API simplifiée pour la création, l’inspection, le débogage et l’analyse des réseaux de neurones dans JAX.

Metrax est disponible en open source et peut être téléchargé sur la documentation officielle.