title: Understanding accumulated gradients in PyTorch
body: <p>I am trying to comprehend inner workings of the gradient accumulation in <code>PyTorch</code>. My question is somewhat related to these two:</p>
  
  <p><a href="https://stackoverflow.com/questions/48001598/why-do-we-need-to-call-zero-grad-in-pytorch#">Why do we need to call zero_grad() in PyTorch?</a></p>
  
  <p><a href="https://stackoverflow.com/questions/44732217/why-do-we-need-to-explicitly-call-zero-grad?noredirect=1&amp;lq=1">Why do we need to explicitly call zero_grad()?</a></p>
  
  <p>Comments to the accepted answer to the second question suggest that accumulated gradients can be used if a minibatch is too large to perform a gradient update in a single forward pass, and thus has to be split into multiple sub-batches.</p>
  
  <p>Consider the following toy example:</p>
  
  <pre><code>import numpy as np
  import torch
  
  
  class ExampleLinear(torch.nn.Module):
  
      def __init__(self):
          super().__init__()
          # Initialize the weight at 1
          self.weight = torch.nn.Parameter(torch.Tensor([1]).float(),
                                           requires_grad=True)
  
      def forward(self, x):
          return self.weight * x
  
  
  if __name__ == "__main__":
      # Example 1
      model = ExampleLinear()
  
      # Generate some data
      x = torch.from_numpy(np.array([4, 2])).float()
      y = 2 * x
  
      optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
  
      y_hat = model(x)          # forward pass
  
      loss = (y - y_hat) ** 2
      loss = loss.mean()        # MSE loss
  
      loss.backward()           # backward pass
  
      optimizer.step()          # weight update
  
      print(model.weight.grad)  # tensor([-20.])
      print(model.weight)       # tensor([1.2000]
  </code></pre>
  
  <p>Which is exactly the result one would expect. Now assume that we want to process the dataset sample-by-sample utilizing gradient accumulation:</p>
  
  <pre><code>    # Example 2: MSE sample-by-sample
      model2 = ExampleLinear()
      optimizer = torch.optim.SGD(model2.parameters(), lr=0.01)
  
      # Compute loss sample-by-sample, then average it over all samples
      loss = []
      for k in range(len(y)):
          y_hat = model2(x[k])
          loss.append((y[k] - y_hat) ** 2)
      loss = sum(loss) / len(y)
  
      loss.backward()     # backward pass
      optimizer.step()    # weight update
  
      print(model2.weight.grad)  # tensor([-20.])
      print(model2.weight)       # tensor([1.2000]
  </code></pre>
  
  <p>Again as expected, the gradient is calculated when the <code>.backward()</code> method is called.</p>
  
  <p>Finally to my question: what exactly happens 'under the hood'? </p>
  
  <p>My understanding is that the computational graph is dynamically updated going from <code>&lt;PowBackward&gt;</code> to <code>&lt;AddBackward&gt;</code> <code>&lt;DivBackward&gt;</code> operations for the <code>loss</code> variable, and that no information about the data used for each forward pass is retained anywhere except for the <code>loss</code> tensor which can be updated until the backward pass.</p>
  
  <p>Are there any caveats to the reasoning in the above paragraph? Lastly, are there any best practices to follow when using gradient accumulation (i.e. can the approach I use in <em>Example 2</em> backfire somehow)?</p>