Como acelerar modelos PyTorch: guia prático do torch.profiler

PyTorch inclui uma ferramenta integrada torch.profiler para análise de desempenho de modelos. Ela rastreia cada operação em GPU e CPU, ajudando a encontrar gargalos nos cálculos e otimizar o tempo de treinamento.

Por que fazer perfilagem de modelos

Sem perfilagem, o desenvolvedor vê apenas o número final: «a época levou 2 horas». Mas por que 2 horas? Para onde foi o tempo? Talvez o dataloader seja lento? Ou a operação na GPU seja ineficiente? Ou a GPU está ociosa, esperando dados da CPU? Tentar otimizar às cegas é adivinhar. Você muda o tamanho do lote, a velocidade de carregamento, a precisão dos cálculos, mas o resultado não melhora porque você está otimizando o gargalo errado. torch.profiler elimina esse desperdício de tempo: mostra a distribuição exata de tempo e memória por operação.

Como funciona o torch.profiler

A ferramenta rastreia a execução do código no nível dos núcleos CUDA e threads da CPU. Para cada operação — por exemplo, multiplicação de matriz na camada Linear ou convolução em Conv2d — registra o tempo de início, fim e memória usada. Os resultados podem ser exportados em um formato compatível com Chrome DevTools e visualizados como uma linha do tempo. As principais métricas do perfilador são:

• Exclusive time — tempo que a operação levou por si só (sem operações aninhadas)
• Inclusive time — tempo da operação junto com todas as chamadas aninhadas
• Memory — pico de uso de memória da GPU durante a operação
• Sync time — tempo gasto na sincronização entre CPU e GPU (gargalo número um)

Quando torch.profiler revela problemas

Um cenário típico: o modelo treina lentamente, você pensa que precisa de uma GPU mais rápida, mas o perfilador mostra que 60% do tempo da época é espera por dados do dataloader. O problema não é na GPU, mas no carregamento de dados. Solução: aumentar num_workers no DataLoader ou usar pinned memory.

Outro exemplo: o modelo treina em um sistema distribuído (várias GPUs), e o perfilador mostra que muito tempo é gasto sincronizando gradientes entre dispositivos. Este é um sinal para otimizar o gráfico de comunicação ou revisar a estratégia de paralelismo.

Um esclarecimento importante: a perfilagem por si só adiciona overhead. Se você executar o perfilador em todo o ciclo de treinamento, ele pode desacelerar os cálculos em 10-30%. Para medições precisas, é melhor fazer perfilagem apenas do bloco de código de interesse (por exemplo, um único passe forward-backward).

O que isso significa

torch.profiler é o primeiro passo antes de qualquer otimização. Não adivinhe onde está o gargalo: faça perfilagem, veja os resultados e otimize exatamente o que está desacelerando. Isso economizará semanas de experimentos e direcionará os esforços para onde eles realmente trarão resultados. Para engenheiros de ML que desejam acelerar o treinamento de modelos, esta é uma ferramenta fundamental.