For a given model I have a normal loss function that is supposed to train the whole net. And I like to use another loss function that only trains just a single output layer. So I don’t like this auxiliary loss to effect the rest of the networks and just train the output layer that is computed on. How can I code this up in Pytorch?
What I tried so far is below however in practice I did not see loss_aux improving.
criterion_main = nn.L1Loss()
criterion_aux = nn.BCELoss()

optimizer_main = nn.optimizer.Adam(net.parameters())
# aux_out is the layer I like to train with criterion_aux
optimizer_aux = nn.optimizer.Adam(net.aux_out.parameters())

...
# in training loop
optimizer_main.zero_grad()
optimizer_aux.zero_grad()
...
loss_main = criterion_main(out, labels)
loss_aux = criterion_aux(out_aux, labels_aux)

loss_main.backward(retain_graph=True)
optimizer_main.step()

loss_aux.backward()
optimizer_aux.step()