HEALPix Curriculum

config/config_curriculum.yml

@@ -0,0 +1,11 @@
+healpix_curriculum:
+  3: 250
+  4: 250
+  5: 250
+  6: 250
+
+curriculum_streams:
+  3: "./config/streams/era5_1deg/"
+  4: "./config/streams/era5_1deg/"
+  5: "./config/streams/era5_1deg/"
+  6: "./config/streams/era5_1deg/"

packages/common/src/weathergen/common/config.py

@@ -23,6 +23,7 @@
 import yaml.constructor
 import yaml.scanner
 from omegaconf import DictConfig, ListConfig, OmegaConf
+from omegaconf.errors import InterpolationKeyError, InterpolationResolutionError
 from omegaconf.omegaconf import open_dict
 
 from weathergen.common.io import StoreType
@@ -139,25 +140,25 @@ def _strip_interpolation(conf: Config) -> Config:
     """Recursively convert interpolated timedelta/datetime objects to strings."""
     stripped = {}
     if OmegaConf.is_dict(conf):
-        for key in list(conf.keys()):
-            key = str(key)
-            if OmegaConf.is_missing(conf, key):
+        for orig_key in list(conf.keys()):
+            str_key = str(orig_key)
+            if OmegaConf.is_missing(conf, orig_key):
                 val = "???"
-            elif OmegaConf.is_config(conf[key]):
-                val = _strip_interpolation(conf[key])
-            elif key.startswith("_"):
+            elif OmegaConf.is_config(conf[orig_key]):
+                val = _strip_interpolation(conf[orig_key])
+            elif str_key.startswith("_"):
                 continue  # Skip hidden/backup keys
-            elif OmegaConf.is_interpolation(conf, key):
-                raw_key = f"_{key}"
+            elif OmegaConf.is_interpolation(conf, orig_key):
+                raw_key = f"_{str_key}"
                 assert raw_key in conf, (
-                    f"Backup key: {raw_key} expected for interpolated key: {key}"
+                    f"Backup key: {raw_key} expected for interpolated key: {orig_key}"
                 )
                 # Retrieve the value from the backup key (resolves interpolation)
                 val = conf[raw_key]
             else:
-                val = conf[key]
+                val = conf[orig_key]
 
-            stripped[key] = val
+            stripped[str_key] = val
     elif OmegaConf.is_list(conf):
         stripped = [
             _strip_interpolation(item) if OmegaConf.is_config(item) else item for item in conf
@@ -456,20 +457,42 @@ def load_merge_configs(
     with open_dict(base_config):
         base_config.from_run_id = from_run_id
         # streams from an overwrite's streams_directory replace inherited streams
-        if any(o.get("streams_directory") is not None for o in overwrite_configs):
+        if any("streams_directory" in o for o in overwrite_configs):
             base_config.streams = None
     # use OmegaConf.unsafe_merge if too slow
     c = OmegaConf.merge(base_config, private_config, *overwrite_configs)
     assert isinstance(c, Config)
     c = _sanitize_time_keys(c)
 
+    if c.get("healpix_curriculum"):
+        istep = c.get("general", {}).get("istep", 0)
+        cumulative = 0
+        current_hl = None
+        unique_curr = {int(hl): steps for hl, steps in c.healpix_curriculum.items()}
+        for hl in sorted(unique_curr.keys()):
+            cumulative += unique_curr[hl]
+            current_hl = hl
+            if istep < cumulative:
+                break
+        c.healpix_level = current_hl
+
+        if c.get("curriculum_streams"):
+            # Support both integer and string keys in the yaml
+            c.streams_directory = c.curriculum_streams.get(current_hl) or c.curriculum_streams.get(
+                str(current_hl)
+            )
+
     return c
 
 
 def _load_streams_in_config(config: Config) -> Config:
     """If the config contains a streams_directory, loads the streams and returns the config with
     the streams set."""
-    streams_directory = config.get("streams_directory", None)
+    try:
+        streams_directory = config.get("streams_directory", None)
+    except (InterpolationKeyError, InterpolationResolutionError):
+        streams_directory = None
+
     config = config.copy()
     if streams_directory is not None:
         streams_directory = Path(streams_directory)

src/weathergen/run_train.py

@@ -147,7 +147,36 @@ def run_continue(args):
     trainer = Trainer(cf.train_logging)
 
     try:
-        trainer.run(cf, devices, args.from_run_id, args.mini_epoch)
+        from_run_id_iter = args.from_run_id
+        mini_epoch_iter = args.mini_epoch
+        first_run = True
+        istep_override = {}
+        while True:
+            if not first_run:
+                cf = config.load_merge_configs(
+                    args.private_config,
+                    from_run_id_iter,
+                    mini_epoch_iter,
+                    args.base_config,
+                    *args.config,
+                    istep_override,
+                    cli_overwrite,
+                )
+                cf = config.set_run_id(cf, cf.general.run_id, True)
+                cf = Trainer.init_ddp(cf)
+                cf.streams = config.load_streams(Path(cf.streams_directory))
+                trainer = Trainer(cf.train_logging)
+
+            trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
+            first_run = False
+
+            if not trainer.cf.get("_curriculum_exit", False):
+                break
+
+            logger.info("Restarting training for next curriculum stage...")
+            from_run_id_iter = trainer.cf.general.run_id
+            mini_epoch_iter = -1
+            istep_override = {"general": {"istep": trainer.cf.general.istep}}
     except Exception:
         extype, value, tb = sys.exc_info()
         traceback.print_exc()
@@ -188,7 +217,35 @@ def run_train(args):
     trainer = Trainer(cf.train_logging)
 
     try:
-        trainer.run(cf, devices)
+        from_run_id_iter = None
+        mini_epoch_iter = None
+        istep_override = {}
+        while True:
+            if from_run_id_iter is not None:
+                cf = config.load_merge_configs(
+                    args.private_config,
+                    from_run_id_iter,
+                    mini_epoch_iter,
+                    args.base_config,
+                    *args.config,
+                    istep_override,
+                    cli_overwrite,
+                )
+                cf = config.set_run_id(cf, cf.general.run_id, True)
+                cf = Trainer.init_ddp(cf)
+                cf.streams = config.load_streams(Path(cf.streams_directory))
+                trainer = Trainer(cf.train_logging)
+                trainer.run(cf, devices, from_run_id_iter, mini_epoch_iter)
+            else:
+                trainer.run(cf, devices)
+
+            if not trainer.cf.get("_curriculum_exit", False):
+                break
+
+            logger.info("Restarting training for next curriculum stage...")
+            from_run_id_iter = trainer.cf.general.run_id
+            mini_epoch_iter = -1
+            istep_override = {"general": {"istep": trainer.cf.general.istep}}
     except Exception:
         extype, value, tb = sys.exc_info()
         traceback.print_exc()

src/weathergen/train/trainer.py

@@ -358,6 +358,24 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
         if self.cf.general.istep > 0 and is_root():
             logger.info(f"Continuing run with learning rate: {self.lr_scheduler.get_lr()}")
 
+        if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum and is_root():
+            unique_curr = {int(hl): steps for hl, steps in self.cf.healpix_curriculum.items()}
+            max_hl = max(unique_curr.keys())
+            if self.cf.healpix_level < max_hl:
+                cumulative = sum(
+                    steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level
+                )
+                logger.info(
+                    f"Curriculum active: Training HEALPix level {self.cf.healpix_level}. "
+                    f"Next stage will begin at istep {cumulative}. "
+                    f"(Note: Total run length is dictated by num_mini_epochs)"
+                )
+            else:
+                logger.info(
+                    f"Curriculum max level ({max_hl}) reached. "
+                    f"Continuing standard training until num_mini_epochs limit."
+                )
+
         # Instantiate loss calculator modules to compute losses
         self.loss_calculator = LossCalculator(cf, self.training_cfg, TRAIN, device=self.device)
         val_cfg = self.validation_cfg
@@ -407,8 +425,16 @@ def run(self, cf, devices, run_id_contd=None, mini_epoch_contd=None):
                 )
             self.save_model(mini_epoch)
 
+            if self.cf.get("_curriculum_exit", False):
+                if is_root():
+                    logger.info("Curriculum stage completed. Exiting training loop.")
+                break
+
         # log final model
-        self.save_model(self.training_cfg.num_mini_epochs)
+        if self.cf.get("_curriculum_exit", False):
+            self.save_model(-1)
+        else:
+            self.save_model(self.training_cfg.num_mini_epochs)
 
     def validate_before_training(self):
         """
@@ -568,6 +594,22 @@ def train(self, mini_epoch):
 
             self.cf.general.istep += 1
 
+            if hasattr(self.cf, "healpix_curriculum") and self.cf.healpix_curriculum:
+                unique_curr = {int(hl): steps for hl, steps in self.cf.healpix_curriculum.items()}
+                max_hl = max(unique_curr.keys())
+                if self.cf.healpix_level < max_hl:
+                    cumulative = sum(
+                        steps for hl, steps in unique_curr.items() if hl <= self.cf.healpix_level
+                    )
+                    if self.cf.general.istep >= cumulative:
+                        if is_root():
+                            logger.info(
+                                f"Curriculum stage for HEALPix level {self.cf.healpix_level} "
+                                f"finished at istep {self.cf.general.istep}. Exiting early."
+                            )
+                        self.cf._curriculum_exit = True
+                        break
+
         self.dataset.advance()
 
     def validate(self, mini_epoch, mode_cfg, batch_size):

src/weathergen/train/trainer_base.py

@@ -136,6 +136,11 @@ def init_ddp(cf):
                     dist.all_reduce(l_seed, op=torch.distributed.ReduceOp.SUM)
                     cf.data_loader_rng_seed = l_seed.item()
 
+        if dist.is_initialized():
+            rank = dist.get_rank()
+            world_size = dist.get_world_size()
+            local_rank = int(os.environ.get("LOCAL_RANK", os.environ.get("SLURM_LOCALID", "0")))
+
         cf.world_size = world_size
         cf.rank = rank
         cf.local_rank = local_rank