Wat is gradient checkpointing?

Gradient checkpointing is een geheugenoptimalisatietechniek tijdens het trainen van diepe neurale netwerken. Het idee is om niet alle tussenactivaties tijdens de voorwaartse pass op te slaan, maar alleen enkele controlepunten (checkpoints). Tijdens de backpropagatie worden de ontbrekende activaties hersteld door een deel van de voorwaartse pass opnieuw te berekenen.

Dit maakt het mogelijk om het geheugenverbruik aanzienlijk te verminderen, wat vooral belangrijk is bij het trainen van zeer diepe modellen of met beperkte GPU-bronnen. Echter, hierdoor neemt de trainingstijd toe, omdat sommige berekeningen opnieuw moeten worden uitgevoerd.

Voorbeeld van gebruik in PyTorch:

from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # complex deel van het model
    return x * x

output = checkpoint(custom_forward, input_tensor)

Dus, gradient checkpointing is een compromis tussen tijd en geheugen, dat het mogelijk maakt om grote modellen te trainen met beperkte middelen.