Taking the derivative of the output of a Neural Network in the forward method w.r.t. input

Yash_Ramani: <p>Hello. I am trying to create an ODE solver. For this purpose, I wish to take the the derivative of the output of the network w.r.t. to its input in the forward method. Here is my code; note that t and N_0 are both of size (batch_size, 1), and I wish to have N_1 be of this size as well.</p>
<pre><code class="lang-auto">    def forward(self, t):
        N_0 = self.net(t)
        N_0.requires_grad = True
        ones = torch.ones_like(N_0)
        N_1 = torch.autograd.grad(N_0, t, ones, retain_graph=True)[0]
        return t, N_0, N_1  
</code></pre>
<p>but I get the error</p>
<pre><code class="lang-auto">RuntimeError: you can only change requires_grad flags of leaf variables.
</code></pre>
<p>in the line N_0.requires_grad = True. This makes sense to me, but my question is: How do I fix this? I have tried the following:</p>
<pre><code class="lang-auto">    def forward(self, t):
        t.requires_grad = True
        N_0 = self.net(t)
        N = N_0.detach()
        N.requires_grad = True
        ones = torch.ones_like(N)
        N_1 = torch.autograd.grad(N, t, ones, retain_graph=True)[0]
        return t, N_0, N_1  
</code></pre>
<p>but I get the error</p>
<pre><code class="lang-auto">RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.
</code></pre>
<p>Note that I will be using t, N_0, N_1 in my loss function which would be used to learn the weights (i.e. I wish to take the derivative N_1 of N_0 w.r.t. t keeping the weights of the network fixed). Any help would be appreciated! Thanks!</p>

eqy: <p>I believe you need to specify that <code>t</code> requires grad rather than <code>N_0</code> which is not a leaf variable as you have observed:</p>
<pre><code class="lang-auto">import torch

c = torch.nn.Conv2d(1,1,1,1)
i = torch.randn(1,1,1,1, requires_grad=True)
o = c(i).sum()
grad = torch.autograd.grad(o, i)
print(grad)
</code></pre>
<pre><code class="lang-auto">(tensor([[[[-0.2108]]]]),)
</code></pre>

Yash_Ramani: <p>Thanks! As you suggested, I made the following change:</p>
<pre><code class="lang-auto">    def forward(self, t):
        t.requires_grad = True
        N_0 = self.net(t)
        ones = torch.ones_like(N_0)
        N_1 = torch.autograd.grad(N_0, t, ones, retain_graph=True)[0]
        return t, N_0, N_1    
</code></pre>
<p>but now I get the error:</p>
<pre><code class="lang-auto">RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn
</code></pre>
<p>in the line</p>
<pre><code class="lang-auto">N_1 = torch.autograd.grad(N_0, t, ones, retain_graph=True)[0]
</code></pre>

eqy: <p>What is being done in <code>self.net</code>? I cannot reproduce the code as written.</p>

Yash_Ramani: <p>I may not be able to reveal the entire code but <a href="http://self.net" rel="noopener nofollow ugc">self.net</a> consists of the layers of the network.</p>
<p>I have a feeling it may be one particular section of my code that could be causing this error (I may be wrong though). In one of the layers in the network, I am initialising the weights of the linear part (before applying the activations) as follows:</p>
<pre><code class="lang-auto">        with torch.no_grad():
            if self.is_first:
                self.linear.weight.uniform_(-1 / self.in_features, 
                                             1 / self.in_features)
</code></pre>
<p>Could using torch.no_grad() be the issue?</p>

AlphaBetaGamma96: <p>Hi <a class="mention" href="/u/yash_ramani">@Yash_Ramani</a></p>
<p>You could also have a look at torch.func, docs here: <a href="https://pytorch.org/docs/stable/func.html" class="inline-onebox" rel="noopener nofollow ugc">torch.func — PyTorch 2.0 documentation</a></p>
<p>That would allow for you to compute the output and gradients at the same time, but its syntax is a bit different. Here’s an example.</p>
<pre><code class="lang-auto">net = Model(*args, **kwargs)
params = dict(net.named_parameters())

def compute_output(params, inputs):
  outputs = torch.func.functional_call(net, params, inputs)
  return outputs, outputs 

def grad_and_value(params, inputs):
  gradients, values = torch.func.grad(compute_output, has_aux=True)(inputs) #use jacrev/jacfwd for multiple outputs
  return gradients, value

#vectorize the gradient/output function
gradients, outputs = vmap(grad_and_value, in_dims=(None, 0))(params, inputs)
</code></pre>

Yash_Ramani: <p>Hi <a class="mention" href="/u/alphabetagamma96">@AlphaBetaGamma96</a>!</p>
<p>Thank you for your answer. I think I have understood the code you provided and have used it in my context as such:</p>
<pre><code class="lang-auto">    def compute_output(self, inputs):
      outputs = torch.func.functional_call(self.net, self.params, inputs)
      return outputs, outputs 

    def grad_and_value(self, inputs):
      gradients, values = torch.func.jacrev(self.compute_output, has_aux=True)(inputs) #use jacrev/jacfwd for multiple outputs
      return gradients, values
    
    def forward(self, t):
        # t.requires_grad = True
        # N_0 = self.net(t)
        # ones = torch.ones_like(N_0)
        # N_1 = torch.autograd.grad(N_0, t, ones, retain_graph=True)[0]
        gradients, outputs = torch.vmap(self.grad_and_value, in_dims=(None, 0))(t)
        return t, outputs, gradients 
</code></pre>
<p>where</p>
<pre><code class="lang-auto">self.params = dict(self.net.named_parameters())
</code></pre>
<p>However, I get the error</p>
<pre><code class="lang-auto">ValueError: vmap(grad_and_value, in_dims=(None, 0), ...)(&lt;inputs&gt;): in_dims is not compatible with the structure of `inputs`. in_dims has structure TreeSpec(tuple, None, [*,
                       *]) but inputs has structure TreeSpec(tuple, None, [*]).
</code></pre>
<p>in the line</p>
<pre><code class="lang-auto">gradients, outputs = torch.vmap(self.grad_and_value, in_dims=(None, 0))(t)
</code></pre>
<p>I also get the same error when using torch.func.grad</p>

AlphaBetaGamma96: <p>Make sure to have a minimal reproducible example, and check that you have an <code>nn.Module</code> object.</p>