feat/ Optimizers and LR scheduler for segmentation module

[MicheleCattaneo]

Addition of Optimizers and LR Schedulers to the Segmentation Module

---

Description

This PR introduces some more flexibility in the choice of optimizer and learning rate for the SegmentationLightningModule.
The constructor additionally takes:

• A string matching the class name of any pytorch optimizer (from torch.optim) and a dictionary of kwargs
• A learning rate (float)
• A string matching the class name of any pytorch learning rate scheduler (from torch.optim.lr_scheduler) and a dictionary of kwargs
  • The kwargs are both the arguments of the constructor and the arguments needed from lightning to define how to condition the step() function call (for certain schedulers only). For example, ReduceLROnPlateau.step() needs to be conditioned on something like the validation loss.

I implemented this following the Lightning documentation

The code is not tested yet, I just wanted to give an idea of the changes that would be introduced first and see whether there is interest in this feature.
Looking forward to hear your opinion!

[colon3ltocard]

Good idea @MicheleCattaneo,

I think some "simple" things can already be configured thru the cli and the yaml settings file : Lightning-AI/pytorch-lightning#18210

and

https://lightning.ai/docs/pytorch/stable/cli/lightning_cli_advanced_3.html#multiple-optimizers-and-schedulers

Have you tried that ? In other words, couldn't we rely on already existing features of the cli to achieve what you want ? and maybe this requires removing some code from our lightningmodule.

[MicheleCattaneo]

Good idea @MicheleCattaneo,

I think some "simple" things can already be configured thru the cli and the yaml settings file : Lightning-AI/pytorch-lightning#18210

and

https://lightning.ai/docs/pytorch/stable/cli/lightning_cli_advanced_3.html#multiple-optimizers-and-schedulers

Have you tried that ? In other words, couldn't we rely on already existing features of the cli to achieve what you want ? and maybe this requires removing some code from our lightningmodule.

Hi,
Yes, I was actually wondering how the CLI could define optimizer and lr scheduler even if the configure_optimizers function does not do that or simply defines a fixed default optimizer (like the SegmentationLightningModule currently does). I am not familiar with the CLI. Do you know exactly what can or can not be done and how does it work behind the hood?

Personally, I feel like sometimes Lightning does too much "magic" and wish that stuff was more explicitly defined, but that's just my personal taste.

[colon3ltocard]

@MicheleCattaneo I think it is time to sit down and think about our lightningmodule(s) a bit more also including my colleagues. I'll organize a virtual meeting during the summer or worse case scenario in september.

[colon3ltocard]

@MicheleCattaneo if you do the same as we advise in the README.md isn't that enough ? Following discussion with colleagues we think a full project as an example/illustration plus a cookie-cutter template would be ideal. Your thoughts ?

[MicheleCattaneo]

Hi!
Thanks for the feedback.

I gave a look and yes, using the CLI, Lightning can inject optimizers and some learning rate schedulers. However I could not manage to use ReduceLROnPlateau. Maybe the schedulers that require an input for their step(..) methods, are not handled by the CLI.

Something else that I noticed it that self.hparams (the result of the method self.save_hyperparameters() called in the constructor that records all arguments passed to __init__) does not keep track of the optimizer and lr scheduler that were injected. If instead those were explicit arguments of the constructor, those would be kept track of. Perhaps that would be helpful for reproducibility.

It seems to me that at the cost of having a more complicated definition of configure_optimizers, we get everything more explicitly defined and we can keep a similar usage of the CLI. The CLI arguments that define the optimizer and lr scheduler will become arguments of the module's constructor, rather than special arguments that make Lightning inject an instance attribute.

By setting default arguments to the module, we can then maintain the possibility of a very simple module instantiation (i.e the use of the common Adam and no lr scheduler).

What do you think? Do you see other downsides?

Could you also clarify on your comment about the full project and cookie-cutter template? What do you suggest?