I have a decoder model.decoder, which is comprised of a series of Convolutional Batchnorm and ReLU layers. I have a latent vector latent_l, which is a 8 dimensional latent vector, say, has the dimension (1, 8, 1, 1), where 1 is the batch size. I am doing torch.autograd.functional.jacobian(model.decoder, latent_l), which is taking a huge amount of time, is there any fast approximation for this jacobian?
There is jacrev, but I am not sure if that works for this example where we pass a decoder as a whole and compute the jacobian of the decoder with respect to the latent vector.
When I use torch.autograd.functional.jacobian(model.decoder, latent_l, vectorize=True), the memory consumption of the GPU increases drastically, leading to the crashing of the program. Is there any efficient way of doing this using Pytorch?
If the dimensionality of your output is greater than 8 (you should use strategy=“forward-mode”) and in that case, your jacobian column-by-column with forward mode with 8 forward passes, with each forward pass taking slightly longer than a normal forward pass.
Given that, the Jacobian taking 10x ~ 15x the amount of time it takes to compute a single forward seems expected.
If you see it taking much longer than 10x forward, one explanation is that your output dimensionality >> input dimensionality and you are using “reverse-mode”.
the default strategy for is “reverse-mode” btw, so you’d need to explicitly pass in “foward-mode”