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) Is too slow. Thus, I am exploring the finite difference method (Finite difference - Wikipedia), 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)
However, it requires a lot of memory. Is there a way to make it better?