Accumulating gradients for large batch training on single GPU.

accumulate_gradients.py

TARGET_BATCH_SIZE, BATCH_FIT_IN_MEMORY = 256, 32
accumulation_steps = int(TARGET_BATCH_SIZE / BATCH_FIT_IN_MEMORY)
network.zero_grad()                            # Reset gradients tensors
for i, (imgs, labels) in enumerate(dataloader):
    
    preds = network(imgs)                      # Forward pass
    loss = loss_function(preds, labels)        # Compute loss function
    loss = loss / accumulation_steps           # Normalize our loss (if averaged)
    loss.backward()                            # Backward pass
    if (i+1) % accumulation_steps == 0:        # Wait for several backward steps
        optim.step()                           # Perform an optimizer step
        network.zero_grad()                    # Reset gradients tensors