How to compute the finite difference Jacobian matrix

marcodena: <p>Dear community,<br>
I need to compute the (differentiable) Jacobian matrix of y = f(z), where y is a (B, 3, 128, 128) aka a batch of images, and z is a (B, 64) vector.</p>
<p>Computing the Jacobian matrix through the pytorch functionial jacobian (<a href="https://pytorch.org/docs/stable/autograd.html#torch.autograd.functional.jacobian" class="inline-onebox" rel="noopener nofollow ugc">Automatic differentiation package - torch.autograd — PyTorch 2.1 documentation</a>) Is too slow. Thus, I am exploring the finite difference method (<a href="https://en.wikipedia.org/wiki/Finite_difference" class="inline-onebox" rel="noopener nofollow ugc">Finite difference - Wikipedia</a>), which is an approximation of the Jacobian. My implementation for B=1 is:</p>
<pre><code>def get_jacobian(net, z, x):
    eps = torch.rand((x.size(0), ), device=x.device)
    delta = torch.sqrt(eps)

    x = x.view(x.size(0), -1)
    m = x.size(1)
    n = z.size(1)

    J = torch.zeros((m, n), device=x.device)
    I = torch.eye(n, device=x.device)

    for j in range(n):
        J[:, j] = (net(z+delta*I[:, j]).view(x.size(0), -1) - x) / delta

    return J

x_fake = mynetfunction(z)
J = get_jacobian(mynetfunction, z, x_fake)
</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/5/456f80f5b2f0dd265ce76cd11a6ed37431664588.png" data-download-href="https://discuss.pytorch.org/uploads/default/456f80f5b2f0dd265ce76cd11a6ed37431664588" title="Screenshot 2021-02-23 at 14.37.14"><img src="https://discuss.pytorch.org/uploads/default/original/3X/4/5/456f80f5b2f0dd265ce76cd11a6ed37431664588.png" alt="Screenshot 2021-02-23 at 14.37.14" data-base62-sha1="9UfQ9lAKWAyFiRsqwgBBSn5vYcU" width="690" height="195" data-dominant-color="F0F0F0"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2021-02-23 at 14.37.14</span><span class="informations">854×242 17.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, it requires a lot of memory. Is there a way to make it better?</p>

albanD: <p>Hi,</p>
<p>I am not sure you are generating your eps correctly: rand is a uniform in [0, 1], not a normal.<br>
Also the value in the formula is for a variance close to 0.</p>
<p>We actually use this to check our gradients here <a href="https://github.com/pytorch/pytorch/blob/master/torch/autograd/gradcheck.py" class="inline-onebox">pytorch/gradcheck.py at master · pytorch/pytorch · GitHub</a><br>
The main things are:</p>
<ul>
<li>Always use double precision otherwise the precision is just really bad</li>
<li>Use delta=1e-6 for double precision. No need to actually do sampling</li>
<li>You can do centered difference if you need more precision: f(z + e_t * eps/2) - f(z - e_t * eps/2). Even even more precise with an additional point at the center. You can check the wikipedia page for the exact formula you need to use in this case.</li>
<li>Don’t create the full I matrix. Just one vector and write a 1 at the proper index into it.</li>
<li>Disable autograd if you’re using nn.Module to avoid extra allocation with <code>@torch.no_grad()</code>
</li>
</ul>