NVIDIA mostrou pipeline completo de otimização de modelo com FastNAS pruning e fine-tuning

NVIDIA lançou um guia prático que demonstra o ciclo completo de otimização de rede neural em um único notebook Google Colab: do treinamento básico ao pruning estrutural e fine-tuning subsequente. Como exemplo, a empresa usa sua própria biblioteca NVIDIA Model Optimizer, o dataset CIFAR-10 e um modelo ResNet20 para mostrar, usando código real, como reduzir a carga computacional de uma rede sem transformar o processo em um conjunto de scripts disparatados e experimentos manuais. O guia começa com a configuração do ambiente e a preparação de um experimento reproduzível.

nvidia-modelopt, torchvision, torchprofile e dependências auxiliares são instalados, seed aleatório e parâmetros de execução são fixos, e então um modo simplificado, mas funcional, é montado para Colab. Nele, o tamanho do lote é definido como 256, o modelo base é treinado por 20 épocas, e o estágio de fine-tuning após o pruning leva mais 12 épocas. Para acelerar, subconjuntos de CIFAR-10 são usados: 12 mil imagens para treinamento, 2 mil para validação e 4 mil para teste.

Depois disso, os autores definem manualmente o ResNet20 em PyTorch com blocos residuais, inicialização de pesos personalizada e lógica explícita de conexão de atalho — ou seja, eles mostram não uma caixa preta, mas uma arquitetura que pode ser rapidamente adaptada à sua própria tarefa. Ênfase especial é colocada na engenharia de scaffold. Para o treinamento, aumentações padrão são aplicadas, incluindo corte aleatório 32x32 e flip horizontal, enquanto a avaliação usa apenas normalização.

O treinamento em si é construído em SGD com momentum 0.9, weight decay 1e-4 e um agendador de taxa de aprendizado usando cosine decay com warmup. O código tem funções separadas para uma época de treinamento, validação, teste e salvamento do melhor checkpoint por acurácia de validação.

Este é um detalhe importante: NVIDIA demonstra não apenas a técnica de compressão em si, mas um pipeline totalmente reproduzível no qual você pode controlar a qualidade do modelo antes e depois da otimização, em vez de simplesmente executar o pruning como um truque único. O estágio-chave é o pruning FastNAS. No exemplo, um limite de 60 milhões de FLOPs é definido, e a configuração de busca é ajustada para que o número de canais e características permaneça divisível por 16.

A acurácia de validação é usada como função de pontuação, e antes de executar, os autores corrigem separadamente a compatibilidade com torchprofile para contar corretamente FLOPs em Colab. Depois disso, Model Optimizer constrói uma subrede leve, salva-a e permite que a arquitetura otimizada seja restaurada para a próxima etapa. Aqui você pode claramente ver como NVIDIA posiciona Model Optimizer: não apenas como uma biblioteca para pruning, mas como uma camada única para técnicas de otimização de modelo.

No repositório oficial, a empresa descreve como um conjunto de ferramentas para pruning, quantization, distillation, sparsity e outros métodos que podem então ser incorporados em infraestrutura de inferência como TensorRT, TensorRT-LLM ou vLLM. Após encontrar a subrede otimizada, o fine-tuning começa. Primeiro, o modelo podado restaurado passa por uma re-verificação, depois retreina com uma taxa de aprendizado mais suave, e finalmente a acurácia antes do pruning, acurácia após o pruning e acurácia após o fine-tuning são comparadas.

Além disso, o número total de parâmetros, o número de pesos diferentes de zero e o tempo gasto em cada estágio—treinamento de baseline, busca FastNAS e recuperação de qualidade—são calculados. Todos os artefatos principais também são salvos: baseline state dict, checkpoint de busca, modelo podado e versão otimizada final. Para engenheiros de ML praticantes, isso é valioso porque o cenário pode ser repetido em sua própria arquitetura com quase nenhuma alteração e incorporado no processo de preparação de um modelo para inferência mais barata e rápida.

A conclusão principal é que NVIDIA faz da otimização de modelo parte do pipeline ML padrão, não uma tarefa separada no estágio final antes da implantação. Esta abordagem é especialmente importante agora, quando o custo da computação, restrições de latência e requisitos para implantação de modelo influenciam cada vez mais decisões arquiteturais tanto quanto a precisão em si. Este material é útil precisamente por sua lógica aplicada: mostra como fazer a transição de uma rede base densa para uma versão mais eficiente em um processo reproduzível e compreensível que você pode realmente executar até mesmo em Google Colab.