diff --git a/config/config_curriculum.yml b/config/config_curriculum.yml new file mode 100644 index 000000000..27aeb1c52 --- /dev/null +++ b/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/" diff --git a/packages/common/src/weathergen/common/config.py b/packages/common/src/weathergen/common/config.py index ccc543cb1..df844a394 100644 --- a/packages/common/src/weathergen/common/config.py +++ b/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 @@ -463,20 +464,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) diff --git a/src/weathergen/run_train.py b/src/weathergen/run_train.py index 5af219f7c..8b71afa0d 100644 --- a/src/weathergen/run_train.py +++ b/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() diff --git a/src/weathergen/train/trainer.py b/src/weathergen/train/trainer.py index 13c60edc5..df58d9b39 100644 --- a/src/weathergen/train/trainer.py +++ b/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): diff --git a/src/weathergen/train/trainer_base.py b/src/weathergen/train/trainer_base.py index d0359892a..b142d60d2 100644 --- a/src/weathergen/train/trainer_base.py +++ b/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