Trainer

[matschreiner]

This PR introduces a new Trainer class that contains all aspects of configuring the training algorithm within the trainer object. The new trainer inherits from pl.Trainer and implements a factory for the configure_optimizers method, assigning this callback to the model when needed.

Previously, if we wanted to change the batch size due to a new learning rate or optimizer, we had to modify the learning rate on the model object while changing the batch size on the trainer object. With this update, both adjustments are handled within the trainer, isolating the configuration of the training algorithm to the Trainer class.

At this point, scheduler and optimizer configuration is limited to:

python neural_lam.train_model.py --scheduler_config ExponentialLR --optimizer_config ...

which will be using the default values for each class. The infrastructure for feeding configurations in a YAML file or similar is not implemented here in order to limit the scope of the PR.

But it would allow for configuring training like

Copy
datastore:
  kind: mdp
  config_path: danra.datastore.yaml
training:
  optimization:
      optimizer: AdamW
      kwargs:
        lr: 0.001
        beta1: ...
        ...
    scheduling:
      scheduler: ExponentialLR
      kwargs:
        gamma: 0.99
        ...

Type of change

• 🐛 Bug fix (non-breaking change that fixes an issue)
• ✨ New feature (non-breaking change that adds functionality)
• 💥 Breaking change (fix or feature that would cause existing functionality to not work as expected)
• 📖 Documentation (Addition or improvements to documentation)

Checklist before requesting a review

• My branch is up-to-date with the target branch - if not update your fork with the changes from the target branch (use pull with --rebase option if possible).
• I have performed a self-review of my code
• For any new/modified functions/classes I have added docstrings that clearly describe its purpose, expected inputs and returned values
• I have placed in-line comments to clarify the intent of any hard-to-understand passages of my code
• I have updated the README to cover introduced code changes
• I have added tests that prove my fix is effective or that my feature works
• I have given the PR a name that clearly describes the change, written in imperative form (context).
• I have requested a reviewer and an assignee (assignee is responsible for merging). This applies only if you have write access to the repo, otherwise feel free to tag a maintainer to add a reviewer and assignee.

Checklist for reviewers

Each PR comes with its own improvements and flaws. The reviewer should check the following:

• the code is readable
• the code is well tested
• the code is documented (including return types and parameters)
• the code is easy to maintain

Author checklist after completed review

• I have added a line to the CHANGELOG describing this change, in a section
  reflecting type of change (add section where missing):
  • added: when you have added new functionality
  • changed: when default behaviour of the code has been changed
  • fixes: when your contribution fixes a bug

Checklist for assignee

• PR is up to date with the base branch
• the tests pass
• author has added an entry to the changelog (and designated the change as added, changed or fixed)
• Once the PR is ready to be merged, squash commits and merge the PR.

[matschreiner]

These arguments are actually not functional yet and serve as placeholders.
In reality the configurations needs to be dictionaries, eg

    optimizer_config = {
        "optimizer": "Adam",
        "kwargs": {"lr": 0.01},
    }

which I imagine should come from the NeuralLamConfig file

[joeloskarsson]

Trying to organize a bit among PRs, so marking this as a draft for now. If those involved want to pick this back up and review this to get it merged please switch it back to ready for review.

[sadamov]

Closing this one. Author noted the args were placeholders rather than functional, and there were no takers for picking up the Trainer factory design. The post-#208 architecture restructured training in a different direction. Thanks @matschreiner.