[Feature] Learning rate scheduling callback from yaml

[Xmaster6y]

What does this PR do?

It implements a callback LRSchedulerCallback to use torch scheduling classes. For further extension, this callback is integrated within the config group callback. The Callback & CallbackNotifier classes were extended to enable callbacks to save data in the state dict (e.g. learning rate schedulers).

It can be loaded using:

defaults:
  ...
  - callback@callbacks.c1: lr_scheduler
  - _self_

Depending on the class, additional arguments can be provided:

callbacks:
  c1:
    scheduler_class: torch.optim.lr_scheduler.CosineAnnealingLR
    scheduler_params:
      T_max: 500
      eta_min: 0.0

It can also be disabled:

defaults:
  ...
  - callback@callbacks: null
  - _self_

or left empty:

defaults:
  ...
  - _self_

To discuss

Maybe the defaults should be without any callback (can be left empty), yet the current LR scheduling doesn't change much training with basic configs. Maybe we should also save the optimizer within the experiment?

Defaults could also look like that:

scheduler_class: torch.optim.lr_scheduler.CosineAnnealingLR

scheduler_params:
  T_max: 500
  eta_min: 0.0

[matteobettini]

This PR is amazing thanks!

There is a lot tho to discuss at once, would it be possible to shard into three:

• new callback for state dict
• new scheduler callback
• hydra callbacks

Ty

[Xmaster6y]

Sounds good!