How to change optimizer and lr scheduler in the middle of training

[icedpanda]

I need to train a model multi-phases with a pre-trained backbone.

For the first 10 epoch, I want to have the backbone frozen and train the classifier only. After epoch 10, I want to start training the whole network. After a certain point (e.g. 100 epochs), I need to enable certain blocks in the network. In regular PyTorch, I would instantiate a new optimizer adding the backbone params, additional required blocks params that I want to train. Then I’d swap both optimizer and lr_scheduler.

I know I can make a multi trainer.fit(). But, what’s the recommended way to do something like this in PL in a callback like BaseFinetuninig?

Here’s my sample code in a callback

    def freeze_before_training(self, pl_module):
        # Here, we are freezing `backbone`
        self.freeze(pl_module.net.encoder)
        if not pl_module.shared_weights:
            self.freeze(pl_module.net.left_encoder)
            self.freeze(pl_module.net.right_encoder)

    def on_train_start(self, trainer, pl_module) -> None:

        if trainer.current_epoch == self._unfreeze_at_epoch:
            print("unfreeze and add param group...")
            pl_module.net.freeze_backbone(False)
            new_optimizer = optim.Adam(
                filter(
                    lambda p: p.requires_grad,
                    pl_module.net.parameters()),
                lr=pl_module.lr,
                weight_decay=pl_module.weight_decay)
            new_schedulers = optim.lr_scheduler.ReduceLROnPlateau(
                new_optimizer,
                mode="min",
                factor=0.1,
                patience=pl_module.scheduler_patience,
                cooldown=3,
            )
            # not sure if its correct or safe to do this
            trainer.optimizers = [new_optimizer]
            trainer.lr_schedulers = [new_schedulers]
    if not pl_module.shared_weights and current_epoch == self._enable_left_view_at_epoch:
        # do the same process
        # unfreeze, and change opt and scheduler
    if not pl_module.shared_weights and current_epoch == self._enable_right_view_at_epoch:
       # do the same process
       # unfreeze, and change opt and scheduler
   # and more conditions and blocks

[four4fish]

@icedpanda Lightning supports multiple optimizers. You can define multiple optimizers and LR_schedules in LightningModule.configure_optimizers(), for example:

class yourModule(LightningModule):
    def __init_(..):

    def configure_optimizers(self):
        optimizer1 = torch.optim.Adam(params, lr)
        optimizer2 = torch.optim.Adam(params, lr)
        lr_scheduler1 = torch.optim.lr_scheduler.StepLR(optimizer1)
        lr_scheduler2 = torch.optim.lr_scheduler. ExponentialLR(optimizer2)
        return [optimizer1, optimizer2], [lr_scheduler1, lr_scheduler2]

   def training_step(self, batch, batch_idx, optimizer_idx):
        if optimizer_idx == 0:
            # do training_step when freeze
            ...
        if optimizer_idx == 1:
            # do training_step when unfreeze
            ...

[icedpanda]

@four4fish How can I change the optimizer based on num_epoch here? From FINETUNING Callbacks, I can free/unfreeze and add_para based on the epoch, but I am not sure how's your code can change the opt during training_step based on epoch?

Would you mind explaining that in detail, please?

Thanks

[quancs]

@icedpanda Hello, I also want to change the optimizer. Do you find a solution now?

[quancs]

Seems the solution by @four4fish works.

@icedpanda To switch optimizer by epoch, you can use self.trainer.current_epoch

    def training_step(self, batch, batch_idx, optimizer_idx):
        if self.trainer.current_epoch < 2:
            if optimizer_idx == 1:
                return None
        elif optimizer_idx == 0:
            return None

        loss = self(batch).sum()
        self.log("train_loss", loss)
        return {"loss": loss}

[icedpanda]

Hi @quancs, I ended up using multiple trainer.fit() as I have many other operations involved, but you can definitely call multiple optimizers in the training_step as suggested.