O que é o checkpointing de gradiente?

O checkpointing de gradiente é uma técnica de otimização de memória durante o treino de redes neurais profundas. A ideia é não guardar todas as ativações intermediárias durante a passagem direta, mas apenas alguns pontos de controlo (checkpoints). Durante a retropropagação, as ativações em falta são reconstruídas recalculando parte da passagem direta.

Isto permite reduzir significativamente o consumo de memória, o que é especialmente importante ao treinar modelos muito profundos ou com recursos GPU limitados. No entanto, isto aumenta o tempo de treino, pois algumas computações precisam de ser repetidas.

Exemplo de uso em PyTorch:

from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # parte complexa do modelo
    return x * x

output = checkpoint(custom_forward, input_tensor)

Assim, o checkpointing de gradiente é um compromisso entre tempo e memória, permitindo treinar modelos grandes com recursos limitados.