Add support for using custom Environments and Strategies ⚡

[amorehead]

What does this PR do?

Adds the ability for one to employ custom (e.g., Lightning Fabric) Environments (e.g., SlurmEnvironment) and Strategies (e.g., DeepSpeedStrategy) during model training or evaluation.

Before submitting

• Did you make sure title is self-explanatory and the description concisely explains the PR?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you list all the breaking changes introduced by this pull request?
• Did you test your PR locally with pytest command?
• Did you run pre-commit hooks with pre-commit run -a command?

Did you have fun?

Make sure you had fun coding 🙃 ⚡

[codecov-commenter]

Codecov Report

Attention: 29 lines in your changes are missing coverage. Please review.

Comparison is base (bddbc24) 83.24% compared to head (27fd562) 78.60%.

Files	Patch %	Lines
src/__init__.py	37.50%	10 Missing ⚠️
src/train.py	60.86%	9 Missing ⚠️
src/eval.py	50.00%	8 Missing ⚠️
src/models/mnist_module.py	50.00%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #608      +/-   ##
==========================================
- Coverage   83.24%   78.60%   -4.65%     
==========================================
  Files          11       12       +1     
  Lines         376      430      +54     
==========================================
+ Hits          313      338      +25     
- Misses         63       92      +29     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[gil2rok]

@amorehead Can you explain your rationale behind adding these additional configs? Super curious to know your thought process.

Is it related to PyTorch Lighting's use of a "strategy" defined here: https://lightning.ai/docs/pytorch/stable/extensions/strategy.html?

And some example strategies:

[amorehead]

Hey, @gil2rok. Yes, Lightning's strategy class is what I had in mind here. I wanted to make it possible with this template to use any arbitrary (e.g., advanced) training strategy such as DeepSpeed with only a few lines of code changes.

[gil2rok]

Hey, @gil2rok. Yes, Lightning's strategy class is what I had in mind here. I wanted to make it possible with this template to use any arbitrary (e.g., advanced) training strategy such as DeepSpeed with only a few lines of code changes.

Because this pull request has not been accepted, how did you use these advanced strategies? Did you just manually add them the PyTorch Lightning trainer arguments?

[amorehead]

As you can see here, now one can specify an optional strategy for Lightning to use e.g., via OmegaConf YAML config files.

[amorehead]

@gil2rok