Multiple optimizers but only one loss

[selflein]

Hey! I have a question regarding this library. Really, like how it forces me to structure my code better. I encountered one problem I did not know how to solve based on the documentation.

Let's say I have two optimizers for two parts of the network, e.g. my configure_optimizers() looks like this:

    def configure_optimizers(self):
        optimizer_encoder = optim.Adam(self.encoder.parameters(), ...)
        optimizer_decoder = optim.Adam(self.decoder.parameters(), ...)
        return [optimizer_encoder, optimizer_decoder]

now in the training loop I forward pass the encoder, then the decoder and compute my loss based on the output:

    def training_step(self, batch, batch_nb, optimizer_idx):
        inp, gt = ...

        encoding = self.encoder(inp)
        pred = self.decoder(encoding)

        loss = F.mse_loss(pred, gt)
        return {'loss': loss}

Since I have two optimizers I have to respect that this function is called two times with different optimizer_idx however I just have one loss to backprop. How would I go about this?

What have you tried?

I tried something like this

    def training_step(self, batch, batch_nb, optimizer_idx):
        if optimizer_idx == 1:
              return {}
        inp, gt = ...

        encoding = self.encoder(inp)
        pred = self.decoder(encoding)

        loss = F.mse_loss(pred, gt)
        return {'loss': loss}

However, this leads to an error since no loss key is present in trainer.py:1392.

[williamFalcon]

in that case, just pass in both sets of params to a single optimizer

[selflein]

But I explicitly want two different learning rates for different parts of the network. That is not really possible with a single optimizer AFAIK. One possibility could be to scale gradients on the weights for which I want lower learning rate before running the optimizer but that is really not a a clean solution.

[amatsukawa]

It is possible with parameter groups using a single optimizer. Your use-case is actually the example in the docs: https://pytorch.org/docs/stable/optim.html#per-parameter-options

[selflein]

That’s nice. Thank you for hint!