Gradient of scalar output w.r.t. input in batches

username: dpernes
cooked: <p>I have a neural network with scalar output and I want to compute the gradient of the output with respect to the input. I know I can use <code>torch.autograd.grad</code> for this purpose, but it only works when the batch size is one and hence the output is a scalar tensor.</p>
<p>However, to boost the speed, I want to work with mini-batches and then compute the derivative of each y[i] (output) w.r.t. each X[i] (input). How can I achieve that?</p>
<p>Minimal example:</p>
<pre><code class="lang-auto">import torch
from torch.autograd import grad

# batch size = 1
x = torch.tensor([[1.]], requires_grad=True)  # shape (1, 1)
y = x**2  # shape (1, 1)
y_x = grad(y, x)  # this works because y is a scalar

# batch size = 2, first attempt
x = torch.tensor([[1.], [1.]], requires_grad=True)  # shape (2, 1)
y = x**2  # shape (2, 1)
y_x = grad(y, x)  # RuntimeError: grad can be implicitly created only for scalar outputs

# batch size = 2, second attempt
x = torch.tensor([[1.], [1.]], requires_grad=True)  # shape (2, 1)
y = x**2  # shape (2, 1)
y_x = []
for i in range(len(y)):
  y_x.append(grad(y[i], x[i]))  # RuntimeError: One of the differentiated Tensors appears to not have been used in the graph.</code></pre>

username: dpernes
cooked: <p>Nvm, this does the job:</p>
<pre><code class="lang-auto">x = torch.tensor([[1.], [1.]], requires_grad=True)  # shape (2, 1)
y = x**2  # shape (2, 1)
y_x = grad(y, x, grad_outputs=torch.ones_like(y), create_graph=True)
</code></pre>

username: pascal_notsawo
cooked: <p>What you just proposed gives me this error:</p>
<pre><code class="lang-auto">RuntimeError: Trying to backward through the graph a second time (or directly access saved variables after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved variables after calling backward.
</code></pre>
<p>I would have used the <code>grad_outputs</code> parameter to avoid the loop:</p>
<pre><code class="lang-auto">x = torch.tensor([[1.], [1.]], requires_grad=True) 
y = x**2  # shape (2, 1)
y_x = grad(y, x, grad_outputs=torch.ones_like(y)) 

"""
tensor([[2.], [2.]])
"""

y = x*torch.tensor([[3.], [2.]]) 
y_x = grad(y, x, grad_outputs=torch.ones_like(y)) 
"""
tensor([[3.], [2.]])
"""
</code></pre>

username: dpernes
cooked: <p>Yeah, I forgot to add the <code>create_graph=True</code>, sorry about that. I fixed it and I’ve written a cleaner solution.</p>

username: Megh_Bhalerao
cooked: <p>I am not sure if this is working for me, when I do this, my input is <code>(batch_size, *)</code>, and output loss function, is of dims <code>batch_size</code> , and I have a model. I use autograd the following way <code>grad_list = torch.autograd.grad(loss_val, inputs = [p for p in net.parameters() if p.requires_grad], grad_outputs=torch.ones_like(loss_val), create_graph=True)</code></p>
<p>However my grad_list seems to accumulate the gradients across all elements in the minibatch, where the size of the minibatch is <code>batch_size</code>, since I see my grad_list to have no indication of the minibatch size anywhere.</p>