I am building an autoencoder, and I would like to take the latent layer for a regression task (with 2 hidden layers and one output layer).
This means that I have two loss functions, one for the AE and one for the regression. I have a few questions:
- Do you suggest adding up both loss values and backprop?
- If I want to backprop each model with respect to its own loss value, how should I implement that in pytorch? should I call loss1.backward() and loss2.backward() then opt.step() ?
Sum them up, the other way around is not feasible at all as you are trying to optimize 2 different latent spaces
I am optimizing the followings:
- The VAE
- The regression
And the input to the regression is the latent space of the VAE.
I would like to use the regression loss to change the weight of the regression portion of the model, and the VAE loss to change the VAE weights
Then, backpropagate regression loss on those layers. You just have to detach regression loss not to backprop it farther than it is necessary. 2 backwards should be ok
Yes, I wanna know how I can do that in Pytorch
latent_space, output = VAE(input)
classification = classifier(latent_space.deatch())
Basically whatever you pass to that fully connected layer, detach() it, it breaks backprop graph.
If there were an issue while backpropagating, you should see a warning saying certain layer’s graph has already been used.