Efficient computation with multiple grad_output's in autograd.grad

username: tjlane
cooked: <p>I want to compute Jacobian matrices using pytorch’s <code>autograd</code>. Autograd natively computes Jacobian-vector products, so I’d simple like to pass an identity matrix to obtain the full Jacobian (ie, Jv = JI = J).</p>
<p>One wrinkle: I’d like to implement both standard reverse-mode AD computation for the Jacobian, but also a forward-mode version (which should be faster for most of my applications) using the following trick due to Jamie Townshend:</p>
<aside class="onebox whitelistedgeneric">
  <header class="source">
      <a href="https://j-towns.github.io/2017/06/12/A-new-trick.html" target="_blank" title="12:00AM - 12 June 2017" rel="nofollow noopener">deep yearning – 12 Jun 17</a>
  </header>
  <article class="onebox-body">
    <img src="" class="thumbnail" width="20" height="20">

<h3><a href="https://j-towns.github.io/2017/06/12/A-new-trick.html" target="_blank" rel="nofollow noopener">A new trick for calculating Jacobian vector products</a></h3>

<p>If you have any questions about this post please ask on the discussion thread on /r/machinelearning.</p>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>

<p>I’ve actually gotten the former (reverse-mode) working, using a few hits found online and some fairly straightforward work:</p>
<pre><code class="lang-auto">def rev_jacobian(fxn, x, n_outputs, retain_graph):
    """
    the basic idea is to create N copies of the input
    and then ask for each of the N dimensions of the
    output... this allows us to compute J with pytorch's
    jacobian-vector engine
    """

    # expand the input, one copy per output dimension
    n_outputs = int(n_outputs)
    repear_arg = (n_outputs,) + (1,) * len(x.size())
    xr = x.repeat(*repear_arg)
    xr.requires_grad_(True)

    # both y and I are shape (n_outputs, n_outputs)
    #  checking y shape lets us report something meaningful
    y = fxn(xr).view(n_outputs, -1)

    if y.size(1) != n_outputs: 
        raise ValueError('Function `fxn` does not give output '
                         'compatible with `n_outputs`=%d, size '
                         'of fxn(x) : %s' 
                         '' % (n_outputs, y.size(1)))
    I = torch.eye(n_outputs, device=xr.device)

    J = autograd.grad(y, xr,
                      grad_outputs=I,
                      retain_graph=retain_graph,
                      create_graph=True,  # for higher order derivatives
                      )

    return J[0]
</code></pre>
<p>However, for the forward-mode version I can’t seem to get <code>autograd</code> to behave. Somehow, I cannot seem to get it to accept the identity matrix as an argument for <code>grad_outputs</code> and NOT return a sum of the columns of the final Jacobian I’d like. But given the success of the reverse-mode implementation, it seems like it should be easy!!</p>
<p>Here’s the current code, with a commented out version of what I’d like to do:</p>
<pre><code class="lang-auto">def fwd_jacobian(fxn, x, n_outputs, retain_graph):
    """
    This implementation is very similar to the above, but with
    one twist. To implement a forward-mode AD with rev-mode
    calls, we first compute the rev-mode VJP for one vector (v)
    then we call d/dv(VJP) `n_outputs` times, one per basis vector,
    to obtain the Jacobian.
    
    This should be faster if `n_outputs` &gt; "n_inputs"

    References
    ----------
    .[1] https://j-towns.github.io/2017/06/12/A-new-trick.html
         (Thanks to Jamie Townsend for this awesome trick!)
    """

    xd = x.detach().requires_grad_(True)
    n_inputs = int(xd.size(0))

    # first, compute *any* VJP 
    v = torch.ones(n_outputs, device=x.device, requires_grad=True)
    y = fxn(xd.view(1,n_inputs)).view(n_outputs)

    if y.size(0) != n_outputs:
        raise ValueError('Function `fxn` does not give output '
                         'compatible with `n_outputs`=%d, size '
                         'of fxn(x) : %s'
                         '' % (n_outputs, y.size(0)))

    vjp = torch.autograd.grad(y, xd, grad_outputs=v, 
                              create_graph=True,
                              retain_graph=retain_graph)[0]
    assert vjp.shape == (n_inputs,)

    # TODO somehow the repeat trick does not work anymore
    #      now that we have to take derivatives wrt v
    #      so loop over basis vectors and compose jacobian col by col

    I = torch.eye(n_inputs, device=x.device)
    J = []
    for i in range(n_inputs):
        Ji = autograd.grad(vjp, v,
                          grad_outputs=I[i],
                          retain_graph=retain_graph,
                          create_graph=True,  # for higher order derivatives
                          )
        J.append(Ji[0])

    return torch.stack(J).t()
</code></pre>
<p>Looping over each column of the final Jacobian is slow, of course, and wasteful – I’m computing the Jacobian N times, and throwing most of it away each time. Since this is the bottleneck in my code, it would be sweet to do it all in one go <img src="https://discuss.pytorch.org/images/emoji/apple/grin.png?v=9" title=":grin:" class="emoji" alt=":grin:">.</p>
<p>Hoping someone with more knowledge of the guts of autograd can advise on the proper use here. Thanks in advance!!</p>
<p>PS. I won’t post just yet for sake of brevity, but if anyone is interested I can post a test case as well to work from.</p>

username: albanD
cooked: <p>Hi,</p>
<p>The base implementations here for me are from Adam’s gists: <a href="https://gist.github.com/apaszke/226abdf867c4e9d6698bd198f3b45fb7">here</a> is the one with full Jacobian and Hessian. And <a href="https://gist.github.com/apaszke/c7257ac04cb8debb82221764f6d117ad">here</a> is the one for the forward mode trick <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: tjlane
cooked: <p>Thanks for those, but they don’t quite answer the question :). My question is: can I do this without looping over the columns of the Jacobian (which Adam does)?</p>

username: albanD
cooked: <p>I’m afraid you cannot. This is a limitation of automatic differentiation… You can only efficiently do vJ or Jv products.</p>

username: tjlane
cooked: <p>Right. I suppose I was operating under the assumption pytorch’s implementation was flexible enough to accept an identity matrix as “v”. That may be wrong.</p>
<p>So if you compare these two implementations, the first gives significantly faster run times (in my hands) than the second. Factor of 10-100x depending on the case. Perhaps this is just to some skipped CUDA overhead? I am trying to understand this behavior and see if I can exploit the same kind of speedup using the “forward mode trick”.</p>
<pre><code class="lang-auto">def _rev_jacobian(fxn, x, n_outputs, retain_graph=True):
    """
    the basic idea is to create N copies of the input
    and then ask for each of the N dimensions of the
    output... this allows us to compute J with pytorch's
    jacobian-vector engine
    """

    # expand the input, one copy per output dimension
    n_outputs = int(n_outputs)
    repear_arg = (n_outputs,) + (1,) * len(x.size())
    xr = x.repeat(*repear_arg)
    xr.requires_grad_(True)

    # both y and I are shape (n_outputs, n_outputs)
    #  checking y shape lets us report something meaningful
    y = fxn(xr).view(n_outputs, -1)

    if y.size(1) != n_outputs:
        raise ValueError('Function `fxn` does not give output '
                         'compatible with `n_outputs`=%d, size '
                         'of fxn(x) : %s'
                         '' % (n_outputs, y.size(1)))
    I = torch.eye(n_outputs, device=xr.device)

    J = autograd.grad(y, xr,
                      grad_outputs=I,
                      retain_graph=retain_graph,
                      create_graph=True,  # for higher order derivatives
                      )

    return J[0]


def _rev_jacobian_simple(fxn, x, n_outputs, retain_graph=True):

    n_outputs = int(n_outputs)

    xd = x.detach()
    xd.requires_grad_(True)
    n_inputs = int(xd.size(0))

    y = fxn(xd.view(1,n_inputs)).view(n_outputs)

    if y.size(0) != n_outputs:
        raise ValueError('Function `fxn` does not give output '
                         'compatible with `n_outputs`=%d, size '
                         'of fxn(x) : %s'
                         '' % (n_outputs, y.size(0)))
    I = torch.eye(n_outputs, device=xd.device)

    J = torch.zeros([n_outputs, n_inputs], device=xd.device)
    for i in range(n_outputs):
        J[i,:] = autograd.grad(y, xd,
                               grad_outputs=I[i],
                               retain_graph=retain_graph,
                               create_graph=True,  # for higher order derivatives
                               )[0]
    return J
</code></pre>

username: Yaroslav_Bulatov
cooked: <p>What if you use the same repeat trick for forward mode? IE, replicate x and v vector and do the same as in reverse mode?</p>
<p><a href="https://colab.research.google.com/drive/1tcm7Lvdv0krpPdaYHtWe7NA2bDbEQ0uj#scrollTo=-jwISHOycQ8x" class="onebox" target="_blank" rel="nofollow noopener">https://colab.research.google.com/drive/1tcm7Lvdv0krpPdaYHtWe7NA2bDbEQ0uj#scrollTo=-jwISHOycQ8x</a></p>

username: albanD
cooked: <p>That would be faster but consume much more memory. That might be ok in your case though !</p>
<p>You can do the same for the “forward mode trick” as well I guess. But not sure this will make the Jacobian computation faster than this one.</p>

username: tjlane
cooked: <p>Tricks on tricks for the win?</p>
<p><a class="mention" href="/u/yaroslav_bulatov">@Yaroslav_Bulatov</a> and <a class="mention" href="/u/alband">@albanD</a> thanks. In the end I copied Yaroslav – doing the “repeat” trick twice. I was hoping I would not have to re-compute N reverse passes, but it seems unavoidable.</p>
<p>With code very similar to Yaroslav’s I get a ~5x speed boost using the repeat trick.</p>
<p>And consistent with <a class="mention" href="/u/alband">@albanD</a> 's expectations the faster, non-looping implementations do use much more memory. So some tradeoff. For me speed is more important and I have memory to spare, so I will probably roll with that :).</p>
<p>Thanks to both of you guys!</p>

username: Kazuki_Osawa
cooked: <p>Instead of using <code>torch.autograd</code>, I’ve tried computing the Jacobian “manually”.</p>
<pre><code class="lang-auto">def manual_jacobian(x, model):
    # x -&gt; fc1 -&gt; s1
    # s1 -&gt; relu -&gt; a1
    # a1 -&gt; fc2 -&gt; s2 (output)

    register_hooks(model)
    output = model(x)

    n = output.size(0)
    fc1 = model.fc1
    fc2 = model.fc2

    # Jacobian s2 (output) -&gt; a1
    Ja1 = torch.stack([fc2.weight] * n)

    # ReLU derivative a1 -&gt; s1
    d = (fc2.data_input &gt; 0).type(fc2.data_input.dtype)

    # Jacobian s2 (output) -&gt; s1
    Js1 = contract('bij,bj-&gt;bij', Ja1, d)

    # Jacobian s2 (output) -&gt; x
    Jx = contract('bij,jk-&gt;bik', Js1, fc1.weight)

    return Jx
</code></pre>
<p>here is <a href="https://github.com/cybertronai/pytorch-sso/blob/dev-master/tests/test_jacobian.py" rel="nofollow noopener">the example code</a>.</p>
<p>By limiting the type of layer (Linear, ReLU, and Conv2d for the example code above), we can know how to compute the Jacobian efficiently without using <code>torch.autograd</code> nor the repeat trick.<br>
For example, <a href="https://github.com/pytorch/pytorch/blob/master/torch/nn/grad.py#L123" rel="nofollow noopener">we can get the Jacobian for ConvNd by using ConvTransposeNd</a></p>
<p>Benchmark on 1 GPU: my approach (<code>manual mode</code>) vs <code>reverse mode</code> for a Linear NN</p>
<pre><code class="lang-auto">bs: 100
n_inputs: 10
n_outputs: 100
loop: 100
-------------
reverse mode: 0.459s
manual mode: 0.195s
(Jx_rev - Jx_man).max(): 0.0
</code></pre>
<p>and it’s <strong>much faster</strong> when it comes to a large number of outputs</p>
<pre><code class="lang-auto">bs: 100
n_inputs: 10
n_outputs: 1000
loop: 100
-------------
reverse mode: 18.032s
manual mode: 1.884s
(Jx_rev - Jx_man).max(): 0.0
</code></pre>
<p>For a CNN</p>
<pre><code class="lang-auto">bs: 1
input shape: torch.Size([3, 32, 32])
n_outputs: 100
loop: 1
-------------
reverse mode: 0.416s
manual mode: 0.232s
(Jx_rev - Jx_man).max(): 3.725290298461914e-09
</code></pre>
<p>To summarize the pros &amp; cons of  <code>manual mode</code>:</p>
<p>pros</p>
<ul>
<li>it can compute Jacobian <strong>in a single backpropagation</strong> for an arbitrary number of outputs.</li>
<li>compare to the repeat trick, <strong>it doesn’t require any copy of the inputs for forwarding</strong>.</li>
<li>it can evaluate the Jacobian layer by layer, so it’s supposed to work even with a large model, e.g., ResNet-50 for ImageNet classification.</li>
</ul>
<p>cons</p>
<ul>
<li>we need to figure out the way to calculate the Jacobian for each module we need for constructing a network.</li>
<li>we need to detect the type and the order of the modules in the network.</li>
<li>it would be more complicated for a network with skip connections.</li>
</ul>
<p>I am trying to overcome these shortcomings of <code>manual mode</code>, and I hope this approach would help.</p>

username: albanD
cooked: <p>HI,</p>
<p>You can check this package that uses this trick in a slightly more general way: <a href="https://github.com/cybertronai/autograd-lib">https://github.com/cybertronai/autograd-lib</a></p>

username: Kazuki_Osawa
cooked: <p>Thanks!<br>
But the Jacobian computation in autograd-lib includes the loop along the output dimension.<br>
<aside class="onebox githubblob">
  <header class="source">
      <a href="https://github.com/cybertronai/autograd-lib/blob/master/autograd_lib/autograd_lib.py#L122-L123" target="_blank" rel="nofollow noopener">github.com</a>
  </header>
  <article class="onebox-body">
    <h4><a href="https://github.com/cybertronai/autograd-lib/blob/master/autograd_lib/autograd_lib.py#L122-L123" target="_blank" rel="nofollow noopener">cybertronai/autograd-lib/blob/master/autograd_lib/autograd_lib.py#L122-L123</a></h4>
<pre class="onebox"><code class="lang-py"><ol class="start lines" start="122" style="counter-reset: li-counter 121 ;">
<li>for idx in range(o):</li>
<li>    output.backward(torch.stack([id_mat[idx]] * n), retain_graph=(retain_graph or idx &lt; o - 1))</li>
</ol></code></pre>


  </article>
  <div class="onebox-metadata">
    
    
  </div>
  <div style="clear: both"></div>
</aside>
</p>
<pre><code class="lang-auto">bs: 128
n_inputs: 1000
n_outputs: 1000
loop: 1
-------------
autograd-lib: 2.100s
reverse mode: 1.193s
manual mode: 0.238s
</code></pre>

username: tjlane
cooked: <p><a class="mention" href="/u/kazuki_osawa">@Kazuki_Osawa</a> thanks for sharing, that is cool. Right now I am only using 2dConv / 2dTransposedConv / 2dBatchNorm / Linear layers, so perhaps it would not be too much work to assemble the components for manual Jacobian calculations. I have been changing my architecture frequently, so the ability to have things “just work” is quite nice :).</p>
<p>I will keep this in mind though if I need more speed.</p>
<p>Curious if you benchmarked JAX at all against your manual mode? At some point just implementing stuff in numpy and using JAX is going to be simpler… or going “full manual” <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>

username: Kazuki_Osawa
cooked: <p><a class="mention" href="/u/tjlane">@tjlane</a><br>
I haven’t tried JAX, but that comparison sounds very interesting! I will also try JAX and put the results at some point.</p>
<p>Thanks!</p>

username: Kazuki_Osawa
cooked: <p>I tried the <a href="https://jax.readthedocs.io/en/latest/notebooks/autodiff_cookbook.html#Jacobians-and-Hessians-using-jacfwd-and-jacrev" rel="nofollow noopener">Jacobian computation by JAX</a> (<code>jax.jacfwd</code> and <code>jax.jacrev</code>) on 1 GPU (NVIDIA Tesla V100)</p>
<p>I defined the same Linear NN for JAX using <a href="https://github.com/google-research/flax/tree/prerelease" rel="nofollow noopener">Flax</a></p>
<pre><code class="lang-auto">class FlaxNet(flax.nn.Module):

    def apply(self, x):
        s1 = flax.nn.Dense(x, features=100, bias=False)
        a1 = flax.nn.relu(s1)
        s2 = flax.nn.Dense(a1, features=1000, bias=False)

        return s2
</code></pre>
<p>Jacobian by JAX</p>
<pre><code class="lang-auto">def jax_jacobian(x, model, mode='fwd'):
    output = model(x)
    if mode == 'fwd':
        Jx = jax.jacfwd(model)(x)
    else:
        Jx = jax.jacrev(model)(x)

    return Jx
</code></pre>
<p>For simplicity, I just set mini-batch size <code>bs</code>= 1.</p>
<pre><code class="lang-auto">bs: 1
n_inputs: 1000
n_outputs: 1000
loop: 100
-------------
jax.jacfwd: 1.959s
jax.jacrev: 1.240s
reverse mode: 2.564s
manual mode: 0.471s
</code></pre>
<p>As I’m a beginner of JAX, I might not be fully utilizing the potential of JAX, but the <code>manual mode</code> in PyTorch is still the fastest.<br>
(<a href="https://github.com/cybertronai/pytorch-sso/blob/dev-master/tests/test_jacobian.py" rel="nofollow noopener">example code</a> for these runs)</p>

username: bnikolic
cooked: <p>Your JAX runs look slow – I’ve tried you code on an oldish CPU and runs significantly faster. A large overhead is being introduced for some reason. I would first try calling the JAX jit method on the jacobian and saving the resulting function to a variable which is then called.</p>
<p>With JIT and jacrev on my CPU (E5-2680 v3 @ 2.50GHz) I get time of 0.393. NB you should call block_until_ready() on last output to ensure asynchronous execution does not cause misleading results.</p>

username: Kazuki_Osawa
cooked: <p>Thanks for sharing!<br>
Yes, JAX could be much faster with JIT on a CPU. But even with JIT compiler, when the dimension is so large that we can ignore the CPU overhead, the manual mode (<code>torch.man</code>) is faster than JAX.</p>
<p>I have updated my <a href="https://github.com/cybertronai/pytorch-sso/blob/dev-master/tests/test_jacobian.py" rel="nofollow noopener">code</a> to measure the time with <code>jax.jit</code> and <code>jax.vmap</code>.</p>
<pre><code class="lang-auto">jac_fun = jax.jit(jax.jacrev(model))
Jx_fn = jax.jit(jax.vmap(jac_fun, in_axes=(0,)))
</code></pre>
<p>On a GPU (NVIDIA Tesla V100),</p>
<pre><code class="lang-auto">mode: torch.auto
bs: 32
n_inputs: 1000
hidden_ndim: 1000
n_outputs: 1000
n_layers: 3
-------------
loop: 100
device: cuda
torch auto rev: 2.591s
</code></pre>
<pre><code class="lang-auto">mode: torch.man
bs: 32
n_inputs: 1000
hidden_ndim: 1000
n_outputs: 1000
n_layers: 3
-------------
loop: 100
device: cuda
torch manual rev: 0.454s
</code></pre>
<pre><code class="lang-auto">mode: jax
bs: 32
n_inputs: 1000
hidden_ndim: 1000
n_outputs: 1000
n_layers: 3
-------------
loop: 100
device: [GpuDevice(id=0)]
jit(jax.jacrev): 1.024s
</code></pre>
<p>This is not about PyTorch vs JAX, but the FLOPS we need to evaluate the Jacobian.<br>
It is the nature of the auto-grad to evaluate the vector-Jacobian product (vjp) or the Jacobian-vector product (jvp), so you need extra computation compared to the manual-mode.</p>
<p>for the reverse-mode, auto-grad requires extra (M_L)(M_1)(M_0) FLOPS.<br>
M_L: output dimension of NN<br>
M_0: input dimension of NN (the first layer)<br>
M_1: output dimension of the first layer</p>