How to train two optimizer with one loss?

[stdoo]

I have read the GAN Demo, it is for two losses.

Suppose I have two modules A and B. The training_step is:

def training_step(self, batch, batch_idx, optimizer_idx):
        x, y = batch
        a1, a2 = self.A(x[0])
        preds = self.B(x[1], a1, a2)
        loss = loss_fn(preds, y)
        return loss

How to train A and B with different optimizer or learning rate?

[awaelchli]

The best is to use manual optimization:
https://pytorch-lightning.readthedocs.io/en/latest/common/optimizers.html#manual-optimization

Then you have control over each optimizer step.

You can get the optimzers with self.optimizers()
and for backward you simply have to replace loss.backward() with self.manual_backward(loss)
Make sure you use the latest version of Lightning for that for optimal support.

Hope this helps

[stdoo]

Hello @awaelchli,

I have implemented the manual optimization and the running process seems OK. However, I am not sure if the implementation is exactly what I need. The code is like that:

    def __init__(self):
        self.automatic_optimization = False
        
    def forward(self, x):
        a = self.A(x[0])
        preds = self.B(x[1], a)
        return preds

    def training_step(self, batch, batch_idx, optimizer_idx):
        opt_a, opt_b = self.optimizers()
        x, y = batch
        preds = self(x)
        loss = loss_fn(preds, y.float())
        # start optimization
        opt_a.zero_grad()
        opt_b.zero_grad()
        self.manual_backward(loss)
        opt_a.step()
        opt_b.step()
        return loss

    def configure_optimizers(self):
        opt_a = Adam(self.A.parameters(), lr=self.lra)
        opt_b = Adam(self.B.parameters(), lr=self.lrb)
        return opt_a, opt_b

The manual_backward is called after zero_grad() of both optimizers, and then step() is called for each optimizer to update their respective parameters.

[akihironitta]

@stdoo The code looks perfect to me.

[awaelchli]

agree, looks pretty solid <3