How to get max epochs in pl.LightningModule?

I wang to get max epochs in pl.LightningModule to set CosineAnnealingLR. It’s need this.

    def configure_optimizers(self):
        optimizer = optim.Adam(self.parameters(), lr=1e-4, weight_decay=5e-5)
        scheduler = optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.current_epoch, eta_min=1e-6)
        return {"optimizer": optimizer, "lr_scheduler": scheduler}

Or there is another better way???

Use self.trainer.max_epochs instead of self.current_epoch

2 Likes

Hey, at the time configure_optimizers is called, the self.trainer property should have been filled already. Meaning that you could get the maximum number of epoch with self.trainer.max_epochs.

Just be aware that if you have early-stopping enabled, you might not actually run the maximum number of epochs (and there is no way, lightning could now how long you’ll run in advance).

Edit: @pull_request was a bit faster (I forgot to click the “send” button :smiley: )

1 Like