NVIDIA dévoile un pipeline complet d'optimisation de modèle avec FastNAS pruning et fine-tuning

NVIDIA a publié un guide pratique qui déroule le cycle complet d'optimisation d'un réseau de neurones dans un seul notebook Google Colab : de l'entraînement de base au pruning structurel, puis au réentraînement. Comme exemple, l'entreprise utilise sa bibliothèque NVIDIA Model Optimizer, le jeu de données CIFAR-10 et le modèle ResNet20, et montre sur du code réel comment réduire la charge de calcul du réseau sans transformer le processus en un tas de scripts épars et d'expériences manuelles.

Le guide commence par la mise en place de l'environnement et la préparation d'une expérience reproductible. On y installe nvidia-modelopt, torchvision, torchprofile et des dépendances annexes, on fixe la graine aléatoire et les paramètres d'exécution, puis on monte une configuration simplifiée mais fonctionnelle pour Colab. Le batch size est fixé à 256, le modèle de base est entraîné pendant 20 époques et l'étape de fine-tuning après le pruning dure 12 époques de plus.

Pour accélérer, les auteurs utilisent des sous-ensembles de CIFAR-10 : 12 mille images pour l'entraînement, 2 mille pour la validation et 4 mille pour le test. Ils définissent ensuite ResNet20 à la main dans PyTorch, avec des blocs résiduels, une initialisation personnalisée des poids et une logique explicite de connexions shortcut. Ce n'est pas une boîte noire, mais une architecture que l'on peut adapter rapidement à sa propre tâche.

Le guide insiste particulièrement sur l'outillage d'ingénierie. À l'entraînement, il applique des augmentations standard, dont un recadrage aléatoire de 32x32 et un retournement horizontal, tandis que l'évaluation ne conserve que la normalisation. L'entraînement repose sur SGD avec momentum 0.

9, weight decay 1e-4 et un planificateur de taux d'apprentissage en cosine decay avec warmup. Le code comporte des fonctions distinctes pour une époque d'entraînement, la validation, le test et la sauvegarde du meilleur checkpoint selon la précision en validation. Le détail compte : NVIDIA ne démontre pas seulement la technique de compression, mais un pipeline complet et reproductible dans lequel on peut contrôler la qualité du modèle avant et après optimisation, au lieu de lancer le pruning comme une astuce ponctuelle.

L'étape clé est le FastNAS pruning. L'exemple fixe une limite de 60 millions de FLOPs et configure la recherche pour que le nombre de canaux et de caractéristiques reste un multiple de 16. La score function est la précision en validation et, avant le lancement, les auteurs corrigent séparément la compatibilité avec torchprofile afin que les FLOPs soient comptés correctement dans Colab. Model Optimizer construit ensuite un sous-réseau allégé, le sauvegarde et permet de restaurer l'architecture déjà optimisée pour l'étape suivante.

On voit bien ici comment NVIDIA positionne Model Optimizer : pas seulement comme une bibliothèque de pruning, mais comme une couche unique pour les techniques d'optimisation de modèles. Dans le dépôt officiel, l'entreprise le décrit comme une boîte à outils de pruning, quantization, distillation, sparsity et d'autres méthodes, qui peuvent ensuite être intégrées à une infrastructure d'inférence comme TensorRT, TensorRT-LLM ou vLLM.

Une fois le sous-réseau optimisé trouvé, le fine-tuning commence. Le modèle élagué restauré est d'abord réévalué, puis réentraîné avec un learning rate plus doux, et le guide compare enfin la précision avant le pruning, après le pruning et après le réentraînement. Il calcule aussi le nombre total de paramètres, le nombre de poids non nuls et le temps passé à chaque étape : entraînement de base, recherche FastNAS et récupération de la précision. Tous les artefacts clés sont également sauvegardés : le state dict de base, le checkpoint de la recherche, le modèle élagué et la version optimisée finale.

Pour les ingénieurs ML en exercice, l'intérêt est de pouvoir reproduire le scénario presque tel quel sur leur propre architecture et de l'intégrer au processus de préparation d'un modèle pour une inférence moins chère et plus rapide.

La principale conclusion est que NVIDIA fait de l'optimisation de modèle une partie du pipeline ML standard, et non une tâche distincte à l'étape finale avant le déploiement. Cette approche compte particulièrement aujourd'hui, alors que le coût de calcul, les contraintes de latence et les exigences de déploiement pèsent de plus en plus sur les choix d'architecture, autant que la précision elle-même.

Le matériel est utile précisément par sa logique appliquée : il montre comment passer d'un réseau de base dense à une version plus efficace dans un processus reproductible et lisible, qui peut réellement être exécuté même dans Google Colab.