Hello, I am trying to calculate the Jacobian matrix of a GAN using functorch, so I need the derivative of each intermediate layer with respect to the input (z), I am using the generator in this tutorial, for example, I need to calculate the derivative of the 5th layer of the generator with respect to z, to do this I am using the chain rule in the following loop:
grads = []
z = sample_batch(20) #Generates a random vector of size (20,100)
inp = z[:,:,None,None] #Reshaping to match the Convolution layer
index = 5
grad = 1
for i in range(index):
jac_func = lambda x: G.main[i](x)
if i != 0 :
inp = Get_Hidden_Output(z, i-1) #Gets the output of the i-1 layer when the input of the model is z
grad = vmap(jac_func)(inp)*grad
grads.append(grad)
However, while computing, I get an error when reaching to the BatchNorm layers of the model, so I would like to know if it is possible to calculate the such operation directly without using the chain rule, any help would be appreciated, just in case the error I get for the BatchNorm layer is:
ValueError: expected 4D input (got 3D input)