Dear community,
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.
Computing the Jacobian matrix through the pytorch functionial jacobian (Automatic differentiation package - torch.autograd — PyTorch 1.7.1 documentation 6) Is too slow. Thus, I am exploring the finite difference method (Finite difference - Wikipedia 3), which is an approximation of the Jacobian. My implementation for B=1 is:
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)

Screenshot 2021-02-23 at 14.37.14854×242 17.6 KB
However, it requires a lot of memory. Is there a way to make it better?