NVIDIA mostró un pipeline completo de optimización de modelos con FastNAS pruning y fine-tuning

NVIDIA lanzó una guía práctica que demuestra el ciclo completo de optimización de redes neuronales en un único notebook de Google Colab: desde el entrenamiento básico hasta el pruning estructural y el fine-tuning posterior. Como ejemplo, la empresa utiliza su propia biblioteca NVIDIA Model Optimizer, el conjunto de datos CIFAR-10 y un modelo ResNet20 para mostrar, usando código real, cómo reducir la carga computacional de una red sin convertir el proceso en un conjunto de scripts dispares y experimentos manuales. La guía comienza con la configuración del entorno y la preparación de un experimento reproducible.

Se instalan nvidia-modelopt, torchvision, torchprofile y dependencias auxiliares, se fijan la semilla aleatoria y los parámetros de ejecución, y luego se monta un modo simplificado pero funcional para Colab. En él, el tamaño del lote se establece en 256, el modelo base se entrena durante 20 épocas, y la etapa de fine-tuning después del pruning toma otras 12 épocas. Para acelerar, se utilizan subconjuntos de CIFAR-10: 12 mil imágenes para entrenamiento, 2 mil para validación y 4 mil para prueba.

Después de eso, los autores definen manualmente ResNet20 en PyTorch con bloques residuales, inicialización de pesos personalizada y lógica explícita de conexiones de atajo, es decir, muestran no una caja negra, sino una arquitectura que puede adaptarse rápidamente a su propia tarea. Se hace especial énfasis en la ingeniería de andamios. Para el entrenamiento, se aplican aumentaciones estándar, incluyendo recorte aleatorio 32x32 y volteo horizontal, mientras que la evaluación utiliza solo normalización.

El entrenamiento en sí se construye sobre SGD con momentum 0.9, weight decay 1e-4 y un programador de tasa de aprendizaje usando cosine decay con warmup. El código tiene funciones separadas para una época de entrenamiento, validación, prueba y guardado del mejor checkpoint por precisión de validación.

Este es un detalle importante: NVIDIA demuestra no solo la técnica de compresión en sí, sino un pipeline completamente reproducible en el que se puede controlar la calidad del modelo antes y después de la optimización, en lugar de simplemente ejecutar el pruning como un truco único. La etapa clave es el pruning FastNAS. En el ejemplo, se establece un límite de 60 millones de FLOPs, y la configuración de búsqueda se ajusta para que el número de canales y características sea divisible por 16.

Se utiliza la precisión de validación como función de puntuación, y antes de ejecutar, los autores corrigen por separado la compatibilidad con torchprofile para contar correctamente FLOPs en Colab. Después de esto, Model Optimizer construye una subred ligera, la guarda y permite restaurar la arquitectura optimizada para el siguiente paso. Aquí se ve claramente cómo NVIDIA posiciona Model Optimizer: no solo como una biblioteca para pruning, sino como una capa única para técnicas de optimización de modelos.

En el repositorio oficial, la empresa lo describe como un conjunto de herramientas para pruning, quantization, distillation, sparsity y otros métodos que luego se pueden incrustar en infraestructura de inferencia como TensorRT, TensorRT-LLM o vLLM. Después de encontrar la subred optimizada, comienza el fine-tuning. Primero, el modelo podado restaurado pasa una recomprobación, luego se reentrena con una tasa de aprendizaje más suave, y finalmente se comparan la precisión antes del pruning, la precisión después del pruning y la precisión después del fine-tuning.

Además, se calcula el número total de parámetros, la cantidad de pesos distintos de cero y el tiempo dedicado a cada etapa—entrenamiento de baseline, búsqueda FastNAS y recuperación de calidad. También se guardan todos los artefactos clave: baseline state dict, checkpoint de búsqueda, modelo podado y versión optimizada final. Para ingenieros de ML que practican, esto es valioso porque el escenario se puede repetir en su propia arquitectura con casi ningún cambio e incrustarse en el proceso de preparación de un modelo para inferencia más barata y rápida.

La conclusión principal es que NVIDIA convierte la optimización de modelos en parte del pipeline ML estándar, no en una tarea separada en la etapa final antes del despliegue. Este enfoque es especialmente importante ahora, cuando el costo de la computación, las restricciones de latencia y los requisitos para el despliegue de modelos influyen cada vez más en las decisiones arquitectónicas tanto como la precisión en sí. Este material es útil precisamente por su lógica aplicada: muestra cómo pasar de una red base densa a una versión más eficiente en un proceso reproducible y comprensible que realmente se puede ejecutar incluso en Google Colab.