Accumulating Gradients

sachinruk: <p>I want to accumulate the gradients before I do a backward pass. So wondering what the right way of doing it is. According to <a href="https://medium.com/huggingface/training-larger-batches-practical-tips-on-1-gpu-multi-gpu-distributed-setups-ec88c3e51255" rel="nofollow noopener">this article</a> it’s (let’s assume equal batch sizes):</p>
<pre><code>model.zero_grad()                                   # Reset gradients tensors
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # Forward pass
    loss = loss_function(predictions, 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
        optimizer.step()                            # Now we can do an optimizer step
        model.zero_grad()
</code></pre>
<p>whereas I expected it to be:</p>
<pre><code>model.zero_grad()                                   # Reset gradients tensors
loss = 0
for i, (inputs, labels) in enumerate(training_set):
    predictions = model(inputs)                     # Forward pass
    loss += loss_function(predictions, labels)       # Compute loss function                              
    if (i+1) % accumulation_steps == 0:             # Wait for several backward steps
        loss = loss / accumulation_steps            # Normalize our loss (if averaged)
        loss.backward()                             # Backward pass
        optimizer.step()                            # Now we can do an optimizer step
        model.zero_grad()     
        loss = 0  
</code></pre>
<p>where I accumulate the loss and then divide by the accumulation steps to average it.</p>
<p>Secondary question, if I am right, would you expect my method to be quicker considering I only do the backward pass every accumulation steps?</p>
<p>This is a crosspost from <a href="https://stackoverflow.com/questions/53331540/accumulating-gradients" rel="nofollow noopener">SO</a>.</p>

albanD: <p>Hi,</p>
<p>This has been discussed in details in <a href="https://discuss.pytorch.org/t/why-do-we-need-to-set-the-gradients-manually-to-zero-in-pytorch/4903/20?u=alband">this post</a>.<br>
Let me know if you have further questions !</p>

sachinruk: <p>Thanks for this. Could you see if the logic in my understanding below is correct?</p>
<p>Suppose that in my case <code>accumulation_steps = 10</code>. Therefore with my second code block, I would be creating the graph ten times, therefore would require ten times the memory. However, what I’m not 100% sure is, is this 10 times the memory to hold the parameters, or 10 times the parameters plus the intermediate values calculated in each batch?</p>
<p>The second question is, isn’t loss.backward() the most compute intensive step, since this is where gradients are calculated. So provided that I can hold 10 graphs in memory, wouldn’t this option be the fastest?</p>

albanD: <p>Hi,</p>
<p>It’s never 10 times the parameters. It’s only 10times the intermediary states.</p>
<p>The amount of work in both case is the same. The difference will be that you call a single backward, but the graph it backward through is 10 times bigger. So compute wise it will be the same.</p>

Stihl: <p>Hi,<br>
If I use T5 language model, there isn’t any batch-norm layer inside, should I use averaged loss then?<br>
Or I can just use loss.backward() on every iteration and opt.step() opt.zero_grad() when iteration % accumulation_steps == 0 ?</p>