How to save the gradient after each batch (or epoch)?

vishalthengane: <p>I have an MLP model and I want to save the gradient after each iteration and average it at the last. How I can do that?</p>
<p>model:</p>
<pre><code class="lang-auto">class MyModel(torch.nn.Module):

    def __init__(self, layers_size, input_size=784, num_classes=10):

        super(MyModel, self).__init__()

        # create input layer
        self.layers = torch.nn.ModuleList(
                    [torch.nn.Linear(input_size, layers_size[0])])
        
        # stack hidden layer
        self.layers.extend([torch.nn.Linear(layers_size[i - 1], 
                        layers_size[i]) for i in range(1, len(layers_size ))])
        
        # add output layer
        self.layers.append(torch.nn.Linear(layers_size[-1], num_classes))
        
    def forward(self, x):
        # iterate over all the layers
        for i, layer in enumerate(self.layers):
            if i == len(self.layers):
                x = func.softmax(layer(x)) # softmax for last layer
            else:
                x = func.relu(layer(x))
        return x

</code></pre>
<p>Here what I like to do:</p>
<pre><code class="lang-auto">for epoch in epochs
	for batch in batches:
		model.forward(batch)
		compute_gradients;
		save(gradients)
		model.backward()

avarage(gradients)
</code></pre>
<p>Thanks in advance.</p>

ptrblck: <p>Each <code>backward()</code> call will accumulate the gradients in the <code>.grad</code> attribute of the parameters.<br>
You could thus accumulate the gradients in your data loop and calculate the average afterwards by iterating all parameters and dividing the <code>.grad</code>s by the number of steps.<br>
Alternatively you could also use the <code>autograd.grad</code> method and manually accumulate the gradients.</p>

vishalthengane: <p>Thank you <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Here what I did now,</p>
<pre><code class="lang-auto">grad_bank = {}
model = MyModel()
avg_counter = 0

for epoch in epochs:
    for inputs, targets in batches:
        opt.zero_grad()
        outputs = model.batch(inputs)
        loss = criterion(outputs, targets)
        loss.backward()

        for idx, param in enumerate(model.parameters()):
            grad_bank[f"layer_{idx}"] += param.grad.data
            avg_counter += 1

        opt.step()

for key in grad_bank:
    grad_bank[key] = grad_bank[key] / avg_counter

</code></pre>
<p>Is it right? Also, How to use <code>autograd.grad</code> method.</p>
<p>Thanks for help.</p>

ptrblck: <p>The loop looks correct. I would recommend not to use the <code>.data</code> attribute and if necessary wrap the code in a <code>with torch.no_grad()</code> block.</p>

vishalthengane: <p>Thanks, <a class="mention" href="/u/ptrblck">@ptrblck</a>,</p>
<p>Will <code>.data</code> create some problem? Is there something I should know?</p>
<p>Thanks!</p>

ptrblck: <p>Yes, the usage of the <code>.data</code> attribute is not recommended, as it might yield unwanted side effects.<br>
Autograd won’t be able to track this operation and will thus not be able to raise a proper error, if your manipulation is incorrect (e.g. by changing the underlying data while the computation graph used the original tensors).<br>
If you don’t want to track this operation, warp it in the <code>no_grad()</code> guard.</p>

fabiola: <p>In the case we use a loss function whose attribute <code>reduction</code> is equal to <code>'mean'</code>, shouldn’t <code>av_counter</code> be outside the batch loop ?<br>
Also, I don’t understand why the counter is inside the parameters() loop. Why should we divide each gradient by the number of layers in the case of a neural network ?</p>

Pavithree_shetty: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I have similar question, does averaging out the gradient of every batch is a good representation of model parameters? My case is I would like to use the gradient of one model as a reference for further computation in another model. So If i store the gradient after every backward() and average it out in the end. Does this represent gradient of entire model ? ( is it similar to calculating gradient had i passed entire dataset in one batch?)</p>

ptrblck: <aside class="quote no-group" data-username="Pavithree_shetty" data-post="8" data-topic="97839">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/pavithree_shetty/48/49675_2.png" class="avatar"> Pavithree_shetty:</div>
<blockquote>
<p>I have similar question, does averaging out the gradient of every batch is a good representation of model parameters?</p>
</blockquote>
</aside>
<p>No, as the gradient does not represent the parameters but the updates performed by the optimizer on the parameters.</p>
<aside class="quote no-group" data-username="Pavithree_shetty" data-post="8" data-topic="97839">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/pavithree_shetty/48/49675_2.png" class="avatar"> Pavithree_shetty:</div>
<blockquote>
<p>So If i store the gradient after every backward() and average it out in the end. Does this represent gradient of entire model ? ( is it similar to calculating gradient had i passed entire dataset in one batch?)</p>
</blockquote>
</aside>
<p>It depends if you want to update the parameters after each <code>backward()</code> call.<br>
If so, then the average of the gradients will not represent the gradient calculated using the entire dataset as the parameters were updated between each step.<br>
Also, if your model contains e.g. batchnorm layers the normalization will be different in training mode as the batch stats will be used which will be different using the entire dataset vs. small batches.</p>

Pavithree_shetty: <p>Thank you for the quick reply <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<p>I am trying to store the gradients of the entire model. The code is given below:</p>
<pre><code>  for step, batch in enumerate(train_dataloader): 
    outputs = model(**batch)
    loss = outputs.loss
    loss = loss / args.gradient_accumulation_steps
    accelerator.backward(loss)
    progress_bar.update(1)
    progress_bar.set_postfix(loss=round(loss.item(), 3))
    del outputs
    gc.collect()
    torch.cuda.empty_cache()
    
    if (step+1) % args.gradient_accumulation_steps == 0 or (step+1) == len(train_dataloader):
        optimizer.step()
        scheduler.step()
        optimizer.zero_grad()
 reference_gradient = [ p.grad.view(-1) if p.grad is not None else torch.zeros(p.numel()) for n, p in 
 model.named_parameters()]
 reference_gradient = torch.cat(reference_gradient)
</code></pre>
<p>My intension is to store the model parameters of entire model to used it for further calculation in another model. But I have 2 questions here,</p>
<ol>
<li>Here the reference_gradient variable always returns 0,  I understand that this happens because, optimizer.zero_grad() is called after every gradient.accumulation steps, and all the gradients are set to 0.</li>
<li>How can I store the model parameters of the entire model?</li>
</ol>

ptrblck: <ol>
<li>
<p>It seems the <code>.grad</code> attribute might either be <code>None</code> and the gradients are never calculated or more likely you are trying to store the reference gradients after calling <code>optimizer.zero_grad()</code> and are explicitly zeroing out the gradients.</p>
</li>
<li>
<p>You could store the <code>state_dict</code> of the model.</p>
</li>
</ol>

Pavithree_shetty: <p>I tried storing the state_dict of the model  <a class="mention" href="/u/ptrblck">@ptrblck</a></p>
<p>torch.save(unwrapped_model.state_dict(),“test.pt”)</p>
<p>However, on loading the model, and calculating the reference gradient, it has all tensors set to 0</p>
<p>import torch<br>
model = torch.load(“test.pt”)<br>
reference_gradient = [ p.grad.view(-1) if p.grad is not None else torch.zeros(p.numel()) for n, p in model.named_parameters()]<br>
reference_gradient = torch.cat(reference_gradient)</p>
<p>output : tensor([0., 0., 0.,  …, 0., 0., 0.]) Could you please correct me, i might be missing something</p>

ptrblck: <p>The <code>state_dict</code> will contain all registered parameters and buffers, but not the gradients.<br>
If you want to store the gradients, your previous approach should work in creating e.g. a <code>list</code> or <code>dict</code> and store the gradients there. Just make sure you are not zeroing them out before storing.</p>