How to compute Jacobian matrix in PyTorch?

username: saan77
cooked: <p>For one of my tasks, I am required to compute a forward derivative of output (not loss function) w.r.t given input X. Mathematically, It would look like this:<img src="https://discuss.pytorch.org/uploads/default/original/2X/1/1fbdfc1b21473084f8a41ffd9234281c003ccd9d.png" alt="Screenshot%20(199)" data-base62-sha1="4wNQd64kB7wXSmHzmXPGcWKCLPn" width="401" height="69"><br>
Which is essential a Jacobian of the output. It is different from backpropagation in two ways. First, we want derivative of network output not the loss function. Second, It is calculated w.r.t to input X rather than network parameters. I think this can be achieved in Tensorflow using <code>tf.gradients()</code>. How do I perform this op in PyTorch? I am not sure if I can use <code>backward()</code> function here.</p>
<p>Thanks</p>

username: jpeg729
cooked: <p>You can use <a href="http://pytorch.org/docs/0.3.1/autograd.html#torch.autograd.grad">torch.autograd.grad</a> for stuff like this.</p>

username: DavideA
cooked: <p>Hi, I think to this day the only way is to use <a href="http://pytorch.org/docs/master/autograd.html#torch.autograd.grad" rel="nofollow noopener">grad</a> function, but you will need to call it j times (once for each output). Unfortunately this requires many backward propagations and scales terribly with the output space of the function.</p>
<p>I have this exact same need, is there a different way to get the Jacobian of a function?</p>

username: saan77
cooked: <p>I came across a different solution which uses <code>backward</code> function. It’s all about playing with <code>backward</code>'s parameters. More information can be found <a href="https://stackoverflow.com/questions/43451125/pytorch-what-are-the-gradient-arguments/47026836" rel="nofollow noopener">here</a>. I am still looking for a better solution, if it exists.</p>

username: DavideA
cooked: <p>Thank you saan77, but I am still unable to understand how it is possible to get the Jacobian with a single backward pass.</p>
<p>The <code>grad_tensors</code> argument of <a href="http://pytorch.org/docs/master/autograd.html#torch.autograd.backward" rel="nofollow noopener">backward</a> seems to work as a weighting mask for the <code>tensors</code> argument in the thread you posted.</p>
<p>I don’t want the gradient of my tensors to be accumulated at leaf nodes. I want to get the gradient of <em>each</em> of my tensors with respect to leaf nodes. Let’s say I have an image classifier, whose input has shape (batchsize, c, h, w) and its output has shape (batchsize, n_classes), I want the jacobian to have shape (batchsize, c, h, w, n_classes).</p>
<p>Did you manage to get something similar?</p>

username: saan77
cooked: <p>Yes, you need to call it n times, where n is the number of output nodes. I am not sure if there is a way to compute this in a single pass. <code>compute_jacobian</code> function of this <a href="https://github.com/ast0414/adversarial-example/blob/master/craft.py" rel="nofollow noopener">script</a> computes jacobian using <code>backward</code>.</p>

username: svd3
cooked: <p>I have the exact same issue. I need to compute jacobian many times, and it’s terribly slow to have that many backward passes.</p>
<p>The <a href="https://github.com/HIPS/autograd" rel="nofollow noopener">python Autograd library</a> is much better for jacobian. I was thinking if I could do the same with pytorch.</p>
<p>I hope they implement jacobian soon.</p>

username: pochoi
cooked: <p>I guess it is related to “reverse-mode vs forward-mode”. As wikipedia <a href="https://en.wikipedia.org/wiki/Automatic_differentiation" rel="nofollow noopener">Automatic_differentiation</a> states, reverse-mode is more efficient for “tensor input scalar output” while forward-mode is more efficient for “scalar input tensor output”. That’s why machine learning library uses reverse-mode.</p>
<p>Jacobian matrix, however, is about “tensor input tensor output”. Not sure which way would be more efficient <img src="https://discuss.pytorch.org/images/emoji/apple/thinking.png?v=5" title=":thinking:" class="emoji" alt=":thinking:">.</p>

username: sbarratt
cooked: <p>The <a href="https://gist.github.com/sbarratt/37356c46ad1350d4c30aefbd488a4faa" rel="nofollow noopener">following code</a> will do the trick with a single call to backward, taking advantage of when the function takes batched inputs.</p>

username: Konpat_Ta_Preechakul
cooked: <p>Interesting, I think it only works with input vectors, I don’t see a way to extend it to parameter vectors.</p>

username: Suzyahyah
cooked: <p><a href="https://suzyahyah.github.io/calculus/pytorch/2018/07/01/Pytorch-Autograd-Backprop.html" rel="nofollow noopener">This (verbose) post</a> may help to explain how to do the reconstruction.</p>
<p>“Because  <code>.backward()</code>  requires gradient arguments as inputs and performs a matrix multiplication internally to give the output (see eq 4), the way to obtain the Jacobian is by feeding in a gradient input which accounts for that specific row of the Jacobian. This is done by providing a mask for the specific dimension in the gradient vector”</p>

username: RylanSchaeffer
cooked: <p><a class="mention" href="/u/suzyahyah">@Suzyahyah</a> 's linked post is from July 2018. Is there a more recent post on computing the end-to-end Jacobian of a network in PyTorch?</p>

username: bnikolic
cooked: <p>This is the ticket for the  forward-mode AD feature request<br>
<a href="https://github.com/pytorch/pytorch/issues/10223" rel="nofollow noopener">https://github.com/pytorch/pytorch/issues/10223</a> . This would enable more efficient Jacobian calculation.</p>

username: pascal_notsawo
cooked: <p>Here are some functions that can help you.</p>
<pre><code class="lang-auto">import torch

def gradient(y, x, grad_outputs=None):
    """Compute dy/dx @ grad_outputs"""
    if grad_outputs is None:
        grad_outputs = torch.ones_like(y)
    grad = torch.autograd.grad(y, [x], grad_outputs = grad_outputs, create_graph=True)[0]
    return grad

def jacobian(y, x):
    """Compute dy/dx = dy/dx @ grad_outputs; 
    for grad_outputs in [1, 0, ..., 0], [0, 1, 0, ..., 0], ...., [0, ..., 0, 1]"""
    jac = torch.zeros(y.shape[0], x.shape[0]) 
    for i in range(y.shape[0]):
        grad_outputs = torch.zeros_like(y)
        grad_outputs[i] = 1
        jac[i] = gradient(y, x, grad_outputs = grad_outputs)
    return jac

def divergence(y, x):
    div = 0.
    for i in range(y.shape[-1]):
        div += torch.autograd.grad(y[..., i], x, torch.ones_like(y[..., i]), create_graph=True)[0][..., i:i+1]
    return div

def laplace(y, x):
    grad = gradient(y, x)
    return divergence(grad, x)
</code></pre>
<p>Illustration</p>
<pre><code class="lang-auto">x = torch.tensor([1., 2., 3.], requires_grad=True)
w = torch.tensor([[1., 2., 3.], [0., 1., -1.]])
b = torch.tensor([1., 2.])
y = torch.matmul(x, w.t()) + b # y = x @ wT + b =&gt; y1 = x1 + 2*x2 + 3*x3 + 1 = 15, y2 = x2 - x3 + 2 = 1
dydx = gradient(y, x)  # =&gt; jacobian(y, x) @ [1, 1]
jac = jacobian(y, x) 
div = divergence(y, x)
```</code></pre>

username: Nathaniel_Merrill
cooked: <p>Hey folks I have some exciting news on this front. I was trying to solve the same problem but for a large network that will not work with batch. Even then the batch method described here is still very slow. Instead I devised a way that is highly efficient for network inputs that are relatively small (in my case a VAE decoder with 32 vector input). All you have to do is use a finite difference jacobian, and then the chainrule becomes completely unnecessary. My implementation is not general so I will not share it here (it’s also in C++ pytorch), but you can find an easy <a href="https://github.com/tobydriscoll/fnc-extras/blob/master/fnc/fdjac.m" rel="noopener nofollow ugc">matlab version of finite difference jacobians that I based mine off of here</a> from my wonderful math professors! Note that I had to tune the delta to work for the neural network since their matlab code is in double prec</p>

username: EigenGuy
cooked: <p>I was searching for a solution for the same problem and found out that Autograd now has a <a href="https://pytorch.org/docs/stable/autograd.html#functional-higher-level-api" rel="noopener nofollow ugc">functional module</a> that solves this exact problem. Specifically, <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian" rel="noopener nofollow ugc"><code>torch.autograd.functional.jacobian</code></a>, given a function and input variables, returns the Jacobian. There are also functions to compute the Hessian, Jacobian-vector-product, etc.</p>

username: Vu_Chau
cooked: <p>Hi,<br>
I implemented the computation of the Jacobian matrix using the <code>torch.autograd.functional.jacobian </code> as below. However, it’s still too slow, please help me improve this.</p>
<pre><code class="lang-auto">def jacobianBatch(f, wrt):
    '''
    Compute the jacobian (derivaties of outputs w.r.t inputs)
    Input:
            f: pytorch model
            wrt: batch of training data
    Output:
            jacobian: J of the batch of training data
    '''
    jacobian = []
    f.eval()
    for i in range(wrt.shape[0]):
        jac = torch.autograd.functional.jacobian(f, wrt[i])
        jacobian.append(jac)
        #
    return torch.stack(jacobian, 0)

</code></pre>

username: Le_ZHENG
cooked: <p>autograd.grad itself is a batch version. So if you use <a class="mention" href="/u/pascal_notsawo">@pascal_notsawo</a> method, you just need to iterate n times (n is your variable dimension rather than the batch dimension). So maybe if your variable dimension is smaller, you can iterate over that.</p>