Stormcast Customization

[jleinonen]

Modulus Pull Request

Description

Modifies the StormCast code to be trainable with custom datasets, as well as adding a few additional features.

Overview of changes:

• Combine .to commands in trainer (device and dtype in one call)
• Save regression and diffusion checkpoints in different directories in rundir to avoid diffusion training trying to load regression checkpoint and vice versa
• List included HRRR channels in dataset config instead of excluded ones
• Standardize HRRR dataloader __getitem__ outputs to "background" and "state"
• Remove retrieval of unused ERA5 time step from HRRR dataloader
• Refactor creation of model inputs and outputs into build_network_condition_and_target, which is used for training, validation and inference
• Support for invariant_array being None, indicating no invariants
• Rename variables in inference.py, move plotting and IO to functions
• Add support for multiple rounds of gradient accumulation to enable large batch sizes
• Add support for training with AMP using float16 or bfloat16
• Align default training parameters with StormCast paper

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.

Dependencies

[jleinonen]

/blossom-ci

[jleinonen]

/blossom-ci

[pzharrington]

LGTM!

[jleinonen]

/blossom-ci