NVIDIA mostra pipeline completo de otimização de modelos com FastNAS pruning e fine-tuning

A NVIDIA publicou um guia prático que percorre o ciclo completo de otimização de uma rede neural em um único notebook do Google Colab: do treino base ao pruning estrutural e ao retreino posterior. Como exemplo, a empresa usa sua biblioteca NVIDIA Model Optimizer, o dataset CIFAR-10 e o modelo ResNet20, e mostra em código real como reduzir a carga computacional da rede sem transformar o processo em um amontoado de scripts soltos e experimentos manuais.

O guia começa com a configuração do ambiente e a preparação de um experimento reproduzível. São instalados nvidia-modelopt, torchvision, torchprofile e dependências auxiliares, fixam-se a semente aleatória e os parâmetros de execução, e então se monta uma configuração simplificada, porém funcional, para o Colab. O batch size é 256, o modelo base treina por 20 épocas e a etapa de fine-tuning após o pruning leva mais 12 épocas.

Para acelerar, são usados subconjuntos do CIFAR-10: 12 mil imagens para treino, 2 mil para validação e 4 mil para teste. Em seguida, os autores definem a ResNet20 manualmente em PyTorch, com blocos residuais, inicialização própria de pesos e lógica explícita de conexões shortcut. Não é uma caixa-preta, mas uma arquitetura que dá para adaptar rapidamente à sua tarefa.

O guia dá destaque especial à estrutura de engenharia. No treino, aplicam-se aumentos de dados padrão, incluindo recorte aleatório de 32x32 e espelhamento horizontal, enquanto na avaliação resta apenas a normalização. O treino usa SGD com momentum 0.9, weight decay 1e-4 e um agendador de taxa de aprendizado com cosine decay e warmup. O código tem funções separadas para uma época de treino, validação, teste e salvamento do melhor checkpoint pela acurácia de validação. O detalhe importa: a NVIDIA demonstra não só a técnica de compressão, mas um pipeline completo e reproduzível, no qual é possível controlar a qualidade do modelo antes e depois da otimização, em vez de simplesmente rodar o pruning como um truque pontual.

A etapa central é o FastNAS pruning. O exemplo define um limite de 60 milhões de FLOPs e configura a busca para que o número de canais e de features continue múltiplo de 16. Como score function, usa-se a acurácia na validação e, antes da execução, os autores corrigem separadamente a compatibilidade com o torchprofile para que os FLOPs sejam contados corretamente no Colab. Depois disso, o Model Optimizer constrói uma sub-rede mais leve, a salva e permite restaurar a arquitetura já otimizada para o passo seguinte.

Aqui fica claro como a NVIDIA posiciona o Model Optimizer: não apenas como biblioteca de pruning, mas como uma camada única para técnicas de otimização de modelos. No repositório oficial, a empresa o descreve como um conjunto de ferramentas para pruning, quantization, distillation, sparsity e outros métodos, que depois podem ser integrados a infraestruturas de inferência como TensorRT, TensorRT-LLM ou vLLM.

Depois de encontrada a sub-rede otimizada, começa o fine-tuning. Primeiro, o modelo podado restaurado passa por nova avaliação; depois, é treinado de novo com um learning rate mais suave; por fim, comparam-se a acurácia antes do pruning, após o pruning e após o retreino. Também são calculados o número total de parâmetros, a quantidade de pesos não nulos e o tempo gasto em cada etapa: treino base, busca com FastNAS e recuperação da qualidade. Todos os artefatos principais também são salvos: o state dict base, o checkpoint da busca, o modelo podado e a versão final otimizada.

Para engenheiros de ML em atividade, o valor está em poder repetir o roteiro quase sem mudanças na própria arquitetura e incorporá-lo ao processo de preparação de um modelo para uma inferência mais barata e rápida.

A principal conclusão é que a NVIDIA faz da otimização de modelos parte do pipeline padrão de ML, e não uma tarefa separada na etapa final antes do deploy. Essa abordagem é especialmente importante agora, quando o custo de computação, as restrições de latência e os requisitos de implantação influenciam cada vez mais as decisões de arquitetura, tanto quanto a própria acurácia.

O material é útil justamente pela lógica aplicada: mostra como sair de uma rede base densa para uma versão mais eficiente em um processo reproduzível e compreensível, que realmente pode ser executado até no Google Colab.