Čo je gradient checkpointing?

Gradient checkpointing je technika optimalizácie pamäte počas tréningu hlbokých neurónových sietí. Myšlienka je taká, že počas priameho prechodu neukladáme všetky medziaktivačné hodnoty, ale len niektoré kontrolné body (checkpoints). Pri spätnom šírení sa chýbajúce aktivácie rekonštruujú opätovným výpočtom časti priameho prechodu.

Toto umožňuje výrazne znížiť spotrebu pamäte, čo je obzvlášť dôležité pri tréningu veľmi hlbokých modelov alebo s obmedzenými GPU zdrojmi. Avšak, tým sa zvyšuje čas tréningu, pretože niektoré výpočty je potrebné opakovať.

Príklad použitia v PyTorch:

from torch.utils.checkpoint import checkpoint

def custom_forward(x):
    # zložitá časť modelu
    return x * x

output = checkpoint(custom_forward, input_tensor)

Gradient checkpointing je teda kompromis medzi časom a pamäťou, ktorý umožňuje trénovať veľké modely s obmedzenými zdrojmi.