Adding additional learning rate scheduler

[wiederm]

Pull Request Summary

This PR will add a number of additional learning rate scheduler and infrastructure that is needed to pass control parameters. Additionally, this PR will add a loss scaling scheduler to dynamically control the weight of each loss component as a function for the current epoch index.

Learning rate scheduling

The default learning rate scheduler used a step function reducing the learning rate whenever a monitored property was not improving for a given number of optimization steps, an alternative is something like the CosineAnnealing learning rate scheduler (and it's variety with warmup/restart) that anneales the learning rate from a starting value to a target value over a specified number of epochs.
Other provided LR scheduler are the OneCycle and Cyclic learning rate scheduler, see the PyTorch documentation for their exact behvior.

Loss component scaling

To prioritize different learning tasks in multi-objective learning runs this PR introduces a linear scaling to each component that can be optionally activated using the keywords target_weights and mixing_step for each component name. This results in a scaling of the component from the weight to the target_weight value using mixing_step as the step size (Note that the sign has to be changed depending on the singe of the slope).

In the training run shown below the force component loss is scaled using from an initial weight of 0.8 to 0.2 using -0.1 step size, and then the training is continued with the target weight:

Key changes

Notable points that this PR has either accomplished or will accomplish.

• CosineAnnealingLR scheduler
• CosineAnnealingWarmRestarts scheduler
• OneCycleLR scheduler
• CyclicLR schedular
• Add loss component scaling scheduler
• Add infrastructure for Pydantic models
• Add tests

Associated Issue(s)

• issue Unsstable multitask learning behavior #271

Pull Request Checklist

• Issue(s) raised/addressed and linked
• Includes appropriate unit test(s)
• Appropriate docstring(s) added/updated
• Appropriate .rst doc file(s) added/updated
• PR is ready for review

[codecov-commenter]

Codecov Report

Attention: Patch coverage is 93.64162% with 11 lines in your changes missing coverage. Please review.

Project coverage is 85.35%. Comparing base (45be449) to head (87c4d3c).
Report is 19 commits behind head on main.

Additional details and impacted files