Gradient accumulation and zeroing in PyTorch

Introduction

Gradient accumulation helps train models with big batches using small memory. Zeroing clears old gradients so new learning is correct.

When your computer memory is too small for large batch training.

When you want to simulate a big batch size by adding small batches.

When training deep neural networks that need stable updates.

When you want to control how often the model updates its knowledge.

When you want to avoid mixing old and new gradient information.

Syntax

PyTorch

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

zero_grad() clears old gradients before new backward pass.

Accumulate gradients over several batches before calling optimizer.step().

Examples

Basic training step: clear gradients, compute gradients, update weights.

PyTorch

optimizer.zero_grad()
loss.backward()
optimizer.step()

Accumulate gradients over 4 batches before updating model.

PyTorch

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

Sample Model

This example shows training a simple linear model on 8 samples. It accumulates gradients over 4 samples before updating weights. We print predictions before and after one epoch to see the change.

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

# Simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(2, 1)
    def forward(self, x):
        return self.linear(x)

# Data: 8 samples, 2 features
inputs = torch.tensor([[1.0, 2.0], [2.0, 1.0], [3.0, 4.0], [4.0, 3.0],
                       [5.0, 6.0], [6.0, 5.0], [7.0, 8.0], [8.0, 7.0]])
labels = torch.tensor([[3.0], [3.0], [7.0], [7.0], [11.0], [11.0], [15.0], [15.0]])

model = SimpleModel()
loss_fn = nn.MSELoss()
optimizer = optim.SGD(model.parameters(), lr=0.01)

accumulation_steps = 4

print('Before training:')
with torch.no_grad():
    preds = model(inputs)
    print(preds.squeeze().tolist())

optimizer.zero_grad()
for i in range(len(inputs)):
    input_i = inputs[i].unsqueeze(0)
    label_i = labels[i].unsqueeze(0)
    output = model(input_i)
    loss = loss_fn(output, label_i)
    loss.backward()
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

print('\nAfter one epoch with gradient accumulation:')
with torch.no_grad():
    preds = model(inputs)
    print(preds.squeeze().tolist())

Important Notes

Always call optimizer.zero_grad() before starting to accumulate gradients.

Gradient accumulation lets you use bigger batch sizes without needing more memory.

Forget to zero gradients and you add old gradients again, which can confuse learning.

Summary

Gradient accumulation adds gradients over multiple batches before updating model.

Zeroing gradients clears old information to keep learning correct.

This technique helps train with limited memory and control update frequency.

Practice

1. What is the main purpose of zeroing gradients in PyTorch during training?

easy

A. To clear old gradients before computing new ones

B. To increase the learning rate automatically

C. To save the model weights to disk

D. To initialize the model parameters

Solution

1. Step 1: Understand gradient accumulation process

   Gradients accumulate by default in PyTorch, so old gradients must be cleared before new backward passes.

2. Step 2: Purpose of zeroing gradients

   Zeroing gradients resets them to zero to avoid mixing old and new gradient information.

3. Final Answer:

   To clear old gradients before computing new ones -> Option A

4. Quick Check:

   Zeroing gradients = clear old gradients [OK]

Hint: Zero gradients before backward to avoid mixing old info [OK]

Common Mistakes:

• Thinking zeroing resets model weights
• Confusing zeroing with saving model
• Believing zeroing changes learning rate

2. Which of the following is the correct way to zero gradients of a model named model in PyTorch?

easy

A. model.zero_grad()

B. model.clear_grad()

C. model.reset_gradients()

D. model.zero_gradients()

Solution

1. Step 1: Recall PyTorch API for zeroing gradients

   The official method to zero gradients is zero_grad() on the model or optimizer.

2. Step 2: Check method names

   Only model.zero_grad() is a valid PyTorch method; others are incorrect or do not exist.

3. Final Answer:

   model.zero_grad() -> Option A

4. Quick Check:

   Correct method name = zero_grad() [OK]

Hint: Use model.zero_grad() to clear gradients [OK]

Common Mistakes:

• Using non-existent methods like clear_grad()
• Confusing zero_grad with reset_gradients
• Adding extra letters to method names

3. Consider this PyTorch code snippet:

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

What is the effect of this code on gradient updates?

medium

A. Updates model every batch without accumulation

B. Updates model only once after all batches

C. Does not update model at all

D. Accumulates gradients over 4 batches before updating

Solution

1. Step 1: Analyze the loop and optimizer calls

   Gradients are zeroed before loop and after every 4 batches. Backward is called every batch, so gradients accumulate for 4 batches.

2. Step 2: Understand optimizer.step() timing

   Optimizer updates weights every 4 batches, using accumulated gradients, then zeroes gradients to start fresh.

3. Final Answer:

   Accumulates gradients over 4 batches before updating -> Option D

4. Quick Check:

   Step every 4 batches = gradient accumulation [OK]

Hint: Step optimizer after several backward calls to accumulate [OK]

Common Mistakes:

• Assuming update happens every batch
• Forgetting gradients accumulate without zeroing
• Thinking zero_grad inside loop clears gradients too early

4. Identify the error in this PyTorch training loop:

for batch in data_loader:
    optimizer.zero_grad()
    outputs = model(batch)
    loss = loss_fn(outputs, targets)
    loss.backward()
    optimizer.step()
    loss.backward()

What is wrong here?

medium

A. Not calling optimizer.zero_grad() before backward

B. Missing optimizer.step() call

C. Calling loss.backward() twice without zeroing gradients

D. Using loss.backward() before computing loss

Solution

1. Step 1: Review order of backward and optimizer calls

   Loss.backward() is called twice without zeroing gradients in between, causing gradient accumulation unintentionally.

2. Step 2: Understand consequences of double backward

   Second backward adds gradients again, which is usually a mistake unless intentional accumulation is desired.

3. Final Answer:

   Calling loss.backward() twice without zeroing gradients -> Option C

4. Quick Check:

   Double backward without zero_grad causes wrong gradients [OK]

Hint: Call backward once per batch unless accumulating intentionally [OK]

Common Mistakes:

• Calling backward multiple times without zeroing
• Forgetting zero_grad before backward
• Misplacing optimizer.step() call

5. You want to train a large model on a GPU with limited memory by accumulating gradients over 8 mini-batches before updating weights. Which sequence correctly implements this in PyTorch?

hard

A. Call optimizer.zero_grad() every batch, call loss.backward() every batch, call optimizer.step() every batch

B. Call optimizer.zero_grad() once before loop, call loss.backward() every batch, call optimizer.step() and zero_grad() every 8 batches

C. Call optimizer.zero_grad() every 8 batches, call loss.backward() every 8 batches, call optimizer.step() every batch

D. Call optimizer.zero_grad() once before loop, call loss.backward() every 8 batches, call optimizer.step() every batch

Solution

1. Step 1: Understand gradient accumulation logic

   Gradients must be zeroed once before accumulation starts, then loss.backward() called every batch to accumulate gradients.

2. Step 2: Timing optimizer step and zero_grad

   After accumulating gradients over 8 batches, call optimizer.step() to update weights, then zero_grad() to reset gradients for next accumulation.

3. Final Answer:

   Call optimizer.zero_grad() once before loop, call loss.backward() every batch, call optimizer.step() and zero_grad() every 8 batches -> Option B

4. Quick Check:

   Zero once, backward every batch, step and zero every 8 batches [OK]

Hint: Zero once, backward each batch, step and zero after N batches [OK]

Common Mistakes:

• Zeroing gradients every batch breaks accumulation
• Calling backward only every N batches misses gradients
• Stepping optimizer too often or too rarely