Computing gradient of output with respect to input -- RuntimeError: Trying to backward through the graph a second time during training

username: muammar
cooked: <p>I am building a model that outputs a scaler, and I need to compute the gradient of those outputs with respect to their inputs. I am doing something like the following:</p>
<pre><code class="lang-auto">

def get_forces(outputs, coordinates):
    forces = []
    for index, (hash, image) in enumerate(coordinates.items()):
        for i, coordinate in image: 
            forces_ = -torch.autograd.grad(
                                      outputs[index], 
                                      coordinate,  
                                      create_graph=True,
                                      retain_graph=True
             )[0]
            forces.append(forces_)
    return torch.stack(forces)


while not converged:
    optimizer.zero_grad()  # clear previous gradients
    outputs = model(inputs)    # Note that inputs have requires_grad = True
    outputs = torch.stack(list(outputs.values())).sum(0)
    

    all_grad = get_gradients(outputs, coordinates[0])

    outputs = {"scalar": outputs, "gradiend": all_grad}

    criterion = CustomLossFunction(gradient=True)
    loss = criterion(outputs, targets)    # This loss uses both scalar and gradient
    
    loss.backward()  
    optimizer.step()

    epoch += 1

    
    if epoch == stop:
        converged = True
</code></pre>
<p>The first epoch passes correctly, but at the second epoch I get:</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time, but 
he saved intermediate results have already been freed. Specify 
retain_graph=True when calling backward the first time.
</code></pre>
<p>How can this be solved? How can I compute the gradient of the output of my neural network with respect to its inputs without having this issue? Meanwhile, I am reading on the forum about how other people have solved this.</p>

username: albanD
cooked: <p>Hi,</p>
<p>This is most likely not linked to your <code>get_gradients()</code> function.<br>
But more to the fact that something that you using in the loop (inputs or coordinates) already has a history (you made some operations with it while it was requiring gradients) and so this part of the graph is shared across iterations.<br>
You want to make sure that this is not the case and that you re-create the graph at each iteration.</p>
<p>In particular here, if you get gradients for coordinates from outputs while only inputs were used. I guess you have a graph that links inputs and coordinates and this one is shared across iterations.<br>
You should move that computation inside the while loop.</p>

username: muammar
cooked: <p>This makes total sense and it is the case. This is what I am doing:</p>
<pre><code class="lang-auto">Coordinates === mapping function ===&gt; Features  ====&gt; model ====&gt; output
</code></pre>
<p>The training loop is effectively from <code>Features</code> instead of <code>Coordinates</code> (these have <code>requires_grad=True</code>). The reason for that is because features are not being optimized and to avoid recomputing them on each epoch I just calculated them once. After your explanation, my <code>mapping function</code> has to be moved into <code>model</code> or the training loop so that inputs are really <code>Coordinates</code> and not <code>Features</code>. Is there any way I can make this work without mapping coordinates into features at every epoch?</p>

username: albanD
cooked: <p>Hi,</p>
<p>There is no magic trick to force the autograd to only keep part of the graph I’m afraid.<br>
If you want to avoid recomputing this, you can call <code>.backward(retain_graph=True)</code> which might increase you peak memory a little bit. Also if you do that, I would recommend you either:</p>
<ul>
<li>Move the content of the while loop into a different function. To make sure that at the end of the loop, all the temporary variable are deleted <em>before</em> starting the next operation</li>
<li>Add a bunch of <code>del foo</code> at the end of your loop for <code>outputs</code>, <code>all_grad</code>, and <code>loss</code>. (again to make sure they don’t stay alive until the next iteration).</li>
</ul>

username: muammar
cooked: <p>Thanks <a class="mention" href="/u/alband">@albanD</a>. Your replies were very helpful. After moving the mapping function to the training loop I can see effectively that <code>.backward()</code> does not complain anymore. Now the next problem I am trying to solve is gradients are basically zero where they should not be <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"> I hope to figure it out.</p>

username: albanD
cooked: <p>You can do <code>.register_hook(fn)</code> on any Tensor and <code>fn</code> will be called with the grad computed for this Tensor. That can help you print out gradients to track it down <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"><br>
Good luck!</p>