Qu'est-ce que la sauvegarde de gradient?

La mise en mémoire tampon du gradient est une technique d'optimisation de la mémoire lors de l'entraînement de réseaux neuronaux profonds. L'idée est de ne pas sauvegarder toutes les activations intermédiaires lors de la passe directe, mais seulement certains points de contrôle (checkpoints). Lors de la rétropropagation, les activations manquantes sont reconstruites en recalculant une partie de la passe directe.

Cela permet de réduire considérablement la consommation de mémoire, ce qui est particulièrement important lors de l'entraînement de modèles très profonds ou avec des ressources GPU limitées. Cependant, cela augmente le temps d'entraînement, car une partie des calculs doit être répétée.

Exemple d'utilisation en PyTorch :

from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # partie complexe du modèle
    return x * x

output = checkpoint(custom_forward, input_tensor)

Ainsi, le checkpointing de gradient est un compromis entre temps et mémoire, permettant d'entraîner de grands modèles avec des ressources limitées.