Training Process

Gradient Accumulation

A useful technique when GPU resources are limited.

1
num_accum = 2
2
optimizer.zero_grad()
3
for epoch in range(10):
4
    running_loss = 0.0
5
    for i, data in enumerate(train_loader, 0):
6
        inputs, labels = data
7
        outputs = net(inputs)
8

9
        loss = criterion(outputs, labels) / num_accum
10
        loss.backward()
11

12
        if i % num_accum == 0:
13
            optimizer.step()
14
            optimizer.zero_grad()

• Model parameters are only updated after num_accum iterations.
• The criterion output is divided by num_accum for normalization.
  • My guess: since the accumulated loss over num_accum steps is applied in a single step, dividing by num_accum gives each individual loss an equal weight, producing a normalizing effect.