Hello, I think I have in principle a relatively simple problem that seems to be fundamentally impossible to solve in pytorch.
Let’s assume we propagate a batch X two times trough the same network E and then calculate a loss. I am then interested in the gradient with respect to the parameters in E but only with respect to the first pass. In other words, the second pass trough E together with the loss function could be considered as the effective total loss function for the first pass of E.
This is also equivalent to the following situation:
We clone (deep copy) E and call it H. Then we propagate X one time trough E and then trough H, that is, out=H(E(X)) and calculate a loss on its output. Finally, we calculate the gradient w.r.t E only.
But obviously I don’t want to clone a model for each batch, so how can I do this more efficiently?
Thanks for helping me out here!
Ps: Here is a use case in practice. Assume you have an autoencoder consisting of an encoder E: x->z and decoder (or generator) G:z->x. Now suppose you don’t only want to impose a reconstruction loss s.t. x=G(E(x))=x_hat but also a latent consistency s.t. z=E(x)=E(G(E(x)))=E(x_hat)=z_hat. As you can see, in order to calculate z_hat you propagate two times trough the encoder E, but you don’t want the encoder to learn how to map x_hat close to z (that would result in a degenerated solution). What you want is to teach the encoder-decoder pair to reconstruct an x_hat, such that for a frozen encoder E the latent variable z_hat=E(x_hat) is close to z=E(x). If I would not freeze the second pass of the Encoder, that would be similar as doing the classic mistake in GAN training, where you update the discriminator as well in the generator training step of the GAN training.
The Problem is, that I can not simply zero_grad() or freeze the parameters of the encoder, because then the encoder does not get any gradients at all.