Cuéntame sobre la acumulación de gradientes. ¿En qué casos se aplica?

La acumulación de gradientes es una técnica utilizada en el entrenamiento de redes neuronales, donde los gradientes se acumulan durante varias mini-batches antes de actualizar los pesos del modelo. Esto permite simular de manera eficiente un tamaño de batch mayor que el que cabe físicamente en la memoria de la GPU.

Se aplica en casos donde, debido a limitaciones de memoria, no es posible usar un tamaño de batch grande, y un tamaño de batch grande es importante para la estabilidad y calidad del entrenamiento. En lugar de actualizar los pesos después de cada mini-batch, los gradientes se suman (acumulan) durante varias iteraciones, y luego se realiza un paso de optimización.

Ejemplo de pseudocódigo:

accumulation_steps = 4
optimizer.zero_grad()
for i, data in enumerate(dataloader):
    outputs = model(data)
    loss = loss_fn(outputs, targets) / accumulation_steps
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

De esta manera, el modelo se actualiza con menos frecuencia, pero con un gradiente equivalente a un batch más grande.