Vectorise Grad and Hessian

Fasc: <p>Hi there,</p>
<p>I am trying to speed up / vectorise the following code:</p>
<pre><code class="lang-auto">def function(input_):

    # Compute prediction for u
    for ii in range(input_.shape[0]):
        x = input_[ii, :]
        x.requires_grad = True
        y = targets[ii]
        u = model(x)

        first_derivative = autograd.grad(u, x, retain_graph=True, create_graph=True)
        hessian_mtrx = hessian(model, x, create_graph=False, strict=False)
</code></pre>
<p>I am confused about the interfaces as the hessian requires “sequences”(<a href="https://pytorch.org/docs/stable/autograd.html" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 2.1 documentation</a>)  but</p>
<pre><code class="lang-auto">x_tup = tuple(x.detach()) # with and without .detach()
hessian_mtrx = hessian(model, x_tup, create_graph=False, strict=False)
</code></pre>
<p>does not work and I get a Typeerror where the batch size is confused for the dimension of the input to the model:</p>
<blockquote>
<p>TypeError: forward() takes 2 positional arguments but 33 were given</p>
</blockquote>
<p>I am similarly confused about the correct and efficient usage of autograd.grad as</p>
<pre><code class="lang-auto">x = input_
x.requires_grad = True
u = model(x)
first_derivative = autograd.grad(u, x, retain_graph=True, create_graph=True)
</code></pre>
<p>gives the following error message:</p>
<blockquote>
<p>RuntimeError: grad can be implicitly created only for scalar outputs</p>
</blockquote>
<p>Any pointers or partial solutions are highly appreciated!</p>
<p>Thanks a lot!</p>
<p>model(x) takes an input of size (n,2) and returns a one-dimensional tensor of length n.</p>

albanD: <p>Hi,</p>
<blockquote>
<p>I am confused about the interfaces as the hessian requires “sequences"</p>
</blockquote>
<p>Not sure what you mean here?<br>
If you function takes a single input, you can just give that to “hessian” and if it takes multiple of them, you can just pass a tuple containing all of them. <code>hessian(model, x)</code> in your case should work fine no?</p>
<blockquote>
<p>RuntimeError: grad can be implicitly created only for scalar outputs</p>
</blockquote>
<p><code>autograd.grad()</code> is just running backward mode AD. And so it will create the <code>grad_outputs</code> for you only if the output, <code>u</code> in your case, is a scalar. In other cases, you will need to provide it explicitly with the value that fits what you want to get.</p>

Fasc: <p>Hi,</p>
<p>thank you for taking your time to look at my problem!</p>
<blockquote>
<p>Not sure what you mean here?<br>
If you function takes a single input, you can just give that to “hessian” and if it takes multiple of them, you can just pass a tuple containing all of them.  <code>hessian(model, x)</code>  in your case should work fine no?</p>
</blockquote>
<p>I assumed that if my network has a 2 dimensional input that a tuple of n 1x2 tensors should work as for the example in the documentation. As for my last comment the batch size n is confused for the input dimension, in this can 2.  Does that make more sense?</p>
<blockquote>
<p><code>autograd.grad()</code>  is just running backward mode AD. And so it will create the  <code>grad_outputs</code>  for you only if the output,  <code>u</code>  in your case, is a scalar. In other cases, you will need to provide it explicitly with the value that fits what you want to get.</p>
</blockquote>
<p>Ok, I think you found the critical point there. I am going to look into this and update this comment with relevant links. Thank you!</p>

keunwoo: <p>I guess the meaning of ‘vectorise’ is to calculate Hessian matrix for each row in input.<br>
If it is true, I had the same problem.<br>
This is the <a href="https://discuss.pytorch.org/t/calculating-diagonal-sum-of-hessian-matrix-for-each-sample/95444">issue</a> that I posted.</p>
<p>I don’t think it is possible to do this efficiently with Pytorch.<br>
So I switched to Jax.<br>
This is the <a href="https://gist.github.com/KeunwooPark/d2b0c94dd88895488060e68b2becb6ce" rel="nofollow noopener">code</a> compares performance of Jax and Pytorch for Hassian diagonal vector calculation task. This might be what you need.</p>
<p>But I’m still figuring out if my approach is correct.</p>

albanD: <aside class="quote no-group" data-username="Fasc" data-post="3" data-topic="95467">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/f/c2a13f/48.png" class="avatar"> Fasc:</div>
<blockquote>
<p>As for my last comment the batch size n is confused for the input dimension, in this can 2. Does that make more sense?</p>
</blockquote>
</aside>
<p>Ho you mean you have a single argument that is a tuple? We can’t really handle that indeed. You can wrap it in a tuple before passing it to the hessian function for this to work: <code>hessian(model, (my_tuple_arg,))</code>.</p>