Computing batch Jacobian efficiently

anshu957: <p>I’m trying to compute Jacobian (and its inverse) of the output of an intermediate layer (block1) with respect to the input to the first layer. The code looks like :</p>
<pre><code class="lang-auto">def getInverseJacobian(net2, x):
    # define jacobian matrix
    # x has shape (n_batches X dim of input vector)
    # Take one input point from x and forward it through 1st block 

    jac = torch.zeros(size=(x.shape[1],x.shape[1]))

    y = net2.block1(x)

    for i in range(x.shape[1]):
        jac[i,:] = torch.autograd.grad(y[0][i],x, create_graph=True)[0]



    # Getting inverse of jacobian using Penrose pseudo-inverse
    jac_inverse = torch.pinverse(jac)

    if torch.isnan(jac_inverse).any():
        print('Nan encountered in Jacobian !')
        sys.exit(0)
    
    return jac_inverse
</code></pre>
<p>This works well for single data in a batch. How do I convert it to make a Jacobian for complete batch without using loop. This function will be called several times in the training and loop would not be ideal.<br>
Any suggestions? Am I calculating Jacobian efficiently in the first place? [It is accurate though]</p>

albanD: <p>Hi,</p>
<p>Yes this looks like the right way to do it.<br>
FYI we now have a built in function that does the same thing: <a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian">https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian</a></p>
<p>There is no better way to compute the jacobian yet I’m afraid. But we’re working on it.</p>

tengerye: <p>Hi, do you mean the in-build function also works on one input point?</p>
<p>By the way, may I ask for an example of <a href="https://torch.autograd.functional.jacobian" rel="noopener nofollow ugc">https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian</a> with respect to the parameters of a network please? I have no clue since the first parameter of <code>jacobian</code> is a function.</p>

albanD: <p>Not sure what you mean by “one input point” could you clarify?</p>
<p>For nn.Module, you can check this answer: <a href="https://discuss.pytorch.org/t/get-gradient-and-jacobian-wrt-the-parameters/98240/3" class="inline-onebox">Get gradient and Jacobian wrt the parameters - #3 by albanD</a></p>

hanshu2018: <p>Hi,</p>
<p>Recently, I met the same problem and tried to do the batch_jacobian operation with for-loops. Although it works, but the runtime is too long. Fianlly,  I implemented the batch_jacobian in another way, it is more efficient and the runtime is close to the tf.GradientTape.</p>
<pre><code class="lang-auto">def batch_jacobian(func, x, create_graph=False):
  # x in shape (Batch, Length)
  def _func_sum(x):
    return func(x).sum(dim=0)
  return autograd.functional.jacobian(_func_sum, x, create_graph=create_graph).permute(1,0,2)
</code></pre>

albanD: <p>Note that if you’re using the latest version of pytorch, there is a <code>vectorize=True</code> flag for <code>functional.jacobian()</code> that might speed things up in some cases <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=10" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

Texfy: <p>I’m unsure what the most efficient implementation is if both my inputs and the outputs are batched. Specifically, if I have <code>inputs</code> of shape <code>[B, n]</code> and <code>func</code> maps to outputs of shape <code>[B, m]</code>, then calling</p>
<pre><code class="lang-python">jac = torch.autograd.functional.jacobian(func, inputs, vectorize=True)
</code></pre>
<p>returns a tensor of shape <code>[B, n, B, m]</code>.</p>
<p>But if there is no interaction between batches, then <code>jac[i, :, j, :]</code> are just zero tensors, and I only really need to compute <code>jac[i, :, :] = jac[i, :, i, :]</code>.</p>
<p>Of course, I can just select the appropriate entries of <code>jac</code>, but I’m wondering if this is still the most efficient approach here since a lot of unnecessary gradients are computed. Is there a better way?</p>
<p>Seems to be the same problem as in <a href="https://discuss.pytorch.org/t/reduce-output-dimensions-of-batch-hessian-without-looping/141087">this thread</a>.</p>

AlphaBetaGamma96: <p>Have a look at <a href="https://github.com/pytorch/functorch" rel="noopener nofollow ugc">functorch</a>, that’ll allow you to vectorize over your batch (so you won’t have those zero tensors).</p>

Texfy: <p>Thank you! That was exactly what I was looking for!</p>
<p>For anyone else wondering, <code>functorch.vmap(functorch.jacref(func))(inputs)</code> does the trick.</p>

AlphaBetaGamma96: <p><code>jacrev</code> not <code>jacref</code>! <img src="https://discuss.pytorch.org/images/emoji/apple/wink.png?v=12" title=":wink:" class="emoji" alt=":wink:" loading="lazy" width="20" height="20"></p>
<p><code>functorch.vmap(functorch.jacrev(func))(inputs)</code> <span class="hashtag">#vectorized</span> reverse-mode AD</p>
<p>For anyone else wondering, it also supports forward-mode AD too via <code>jacfwd</code> if you need that too!</p>

Francisco_Garcia: <p>Thank you so much, I have been trying to solve this problem for a week!</p>

AlphaBetaGamma96: <p>Just an update for anyone who reads the thread in the future, as of PyTorch2,the <code>functorch</code> library is now included in pytorch. So you can replace <code>functorch</code> with <code>torch.func</code>, for the most part the syntax is the same except if you have an <code>nn.Module</code> you’ll need to create a ‘functional’ version of your model.</p>
<p>For example,</p>
<pre><code class="lang-auto">model = myModel(*args, **kwargs) #our network

from torch.func import vmap, jacrev, functional_call

params = dict(model.named_parameters())
inputs = torch.randn(batch_size, input_size) #random input data

def fmodel(params, inputs): #functional version of model
  return functional_call(model, params, inputs)

result = vmap(jacrev(fmodel, argnums=(1)), in_dims=(None,0))(params, inputs)
</code></pre>
<p>The documentation for <code>torch.func</code> can be found <a href="https://pytorch.org/docs/stable/func.html" rel="noopener nofollow ugc">here</a>.</p>
<p>Also, if you’re migrating from <code>functorch</code> to <code>torch.func</code> they have a documentation page on the changes between them <a href="https://pytorch.org/docs/stable/func.migrating.html" rel="noopener nofollow ugc">here</a>.</p>

Josuelmet: <p>I was trying to batch the Jacobian for a StyleGAN2 model, which was not directly compatible with functorch or vectorization for some reason (probably some custom hooks somewhere). Doing _func_sum really helps!</p>