feat(weather): Improve configuration for DLWP Healpix and GraphCast examples

- Added configurable checkpoint directory to DLWP Healpix config and training script.
- Implemented Trainer logic to use specific checkpoint directory.
- Updated utils.py to respect exact checkpoint path.
- Made Weights & Biases entity and project configurable in GraphCast example.

Author: dran-dev

CHANGELOG.md

@@ -71,6 +71,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Support passing custom tokenizer, detokenizer, and attention `Module`s in
   experimental DiT architecture
 - Improved Transolver training recipe's configuration for checkpointing and normalization.
+- Improved configuration for DLWP Healpix (checkpoint directory) and GraphCast (W&B settings).
 - Bumped `multi-storage-client` version to 0.33.0 with rust client.
 
 ### Fixed

examples/weather/dlwp_healpix/configs/config.yaml

@@ -22,6 +22,7 @@ defaults:
 
 experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S}
 output_dir: outputs/${experiment_name}
+checkpoint_dir: ${output_dir}/tensorboard/checkpoints
 # checkpoints names are in the form training-state-<name>.mdlus
 checkpoint_name: last
 load_weights_only: false

examples/weather/dlwp_healpix/configs/config_hpx32_coupled_dlom.yaml

@@ -22,6 +22,7 @@ defaults:
 
 experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S}
 output_dir: outputs/${experiment_name}
+checkpoint_dir: ${output_dir}/tensorboard/checkpoints
 # checkpoints names are in the form training-state-<name>.mdlus
 checkpoint_name: last
 load_weights_only: false

examples/weather/dlwp_healpix/configs/config_hpx32_coupled_dlwp.yaml

@@ -22,6 +22,7 @@ defaults:
 
 experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S}
 output_dir: outputs/${experiment_name}
+checkpoint_dir: ${output_dir}/tensorboard/checkpoints
 # checkpoints names are in the form training-state-<name>.mdlus
 checkpoint_name: last
 load_weights_only: false

examples/weather/dlwp_healpix/configs/trainer/default.yaml

@@ -26,3 +26,4 @@ early_stopping_patience: null
 amp_mode: "fp16"
 graph_mode: "train_eval"
 output_dir: ${output_dir}
+checkpoint_dir: ${checkpoint_dir}

examples/weather/dlwp_healpix/configs/trainer/dlom.yaml

@@ -26,4 +26,5 @@ early_stopping_patience: null
 amp_mode: "fp16"
 graph_mode: "train_eval"
 output_dir: ${output_dir}
+checkpoint_dir: ${checkpoint_dir}
 max_norm: 0.25

examples/weather/dlwp_healpix/configs/trainer/dlwp.yaml

@@ -26,4 +26,5 @@ early_stopping_patience: null
 amp_mode: "fp16"
 graph_mode: "train_eval"
 output_dir: ${output_dir}
+checkpoint_dir: ${checkpoint_dir}
 max_norm: 0.25

examples/weather/dlwp_healpix/train.py

@@ -99,18 +99,23 @@ def train(cfg):
     iteration = 0
     epochs_since_improved = 0
 
+    # Determine checkpoint directory
+    if cfg.get("checkpoint_dir") is None:
+        # Fallback to default structure if not specified in config
+        cfg.checkpoint_dir = str(
+            Path(cfg.get("output_dir"), "tensorboard", "checkpoints")
+        )
+
+    checkpoint_dir = Path(cfg.checkpoint_dir)
+
     # Prepare training under consideration of checkpoint if given
     if cfg.get("checkpoint_name", None) is not None:
         checkpoint_path = Path(
-            cfg.get("output_dir"),
-            "tensorboard",
-            "checkpoints",
+            checkpoint_dir,
             "training-state-" + cfg.get("checkpoint_name") + ".mdlus",
         )
         optimizer_path = Path(
-            cfg.get("output_dir"),
-            "tensorboard",
-            "checkpoints",
+            checkpoint_dir,
             "optimizer-state-" + cfg.get("checkpoint_name") + ".ckpt",
         )
         if checkpoint_path.exists():

examples/weather/dlwp_healpix/trainer.py

@@ -53,6 +53,7 @@ def __init__(
         device: torch.device = torch.device("cpu"),
         output_dir: str = "/outputs/",
         max_norm: float = None,
+        checkpoint_dir: str = None,
     ):
         """
         Constructor.
@@ -105,6 +106,10 @@ def __init__(
             num_shards=self.dist.world_size, shard_id=self.dist.rank
         )
         self.output_dir_tb = os.path.join(output_dir, "tensorboard")
+        if checkpoint_dir is None:
+            self.checkpoint_dir = os.path.join(self.output_dir_tb, "checkpoints")
+        else:
+            self.checkpoint_dir = checkpoint_dir
 
         # set the other parameters
         self.optimizer = optimizer
@@ -524,7 +529,7 @@ def fit(
                         iteration,
                         validation_error,
                         epochs_since_improved,
-                        self.output_dir_tb,
+                        self.checkpoint_dir,
                     ),
                 )
                 thread.start()

examples/weather/dlwp_healpix/utils.py

@@ -51,10 +51,7 @@ def write_checkpoint(
     :param dst_path: Path where the checkpoint is written to
     :param keep_n_checkpoints: Number of best checkpoints that will be saved (worse checkpoints are overwritten)
     """
-    root_path = os.path.join(
-        dst_path,
-        "checkpoints",
-    )
+    root_path = dst_path
     # root_path = os.path.dirname(ckpt_dst_path)
     ckpt_dst_path = os.path.join(
         root_path,