Sgrasse/develop/issue 106 (ecmwf#159)

* update dependencies to include 'omegaconf'

* rewrite deserialization of streams config for omegaconf

* move deserialization of stream configs

* fix streams deserialization

* rewrite deserialization of private config for omegaconf

* rewrite deserialization of overwrite config for omegaconf

* add method to deserialize default config

* consistent names for deserialization methods

* transform methods of Config object into functions

* Change implementation of config object to OmegaConf

* add method to load/create new omega config

* formatting

* fix: create "result" directory if not already present

* fix: set 'Path' as default factory in CLI

Using 'Path' as default factory in CLI instead of 'str' avoids the 'str'-'Path' cast potentially casting 'None' to 'Path' (which fails).
The cast would be needed to comply with the 'Path' | 'None' interface of subsequent methods.

* fix: correctly call 'OmegaConf.to_container'

* fix: change return type of stream deserialization to comply with usage

* remove unused private_conf from call to trainer

* load default config from file instead of setting in code

* make evaluate / train continue compatible

* fix: use -2 to select None as epoch value

* remove redundant options that dont differ from default values

* change load_config to account for evaluation/continue

* use 'load_config' for evaluation/continue

* autoformatting

* reenable --config flag for training

* Small improvements suggested by review.

- docstrings
- pure stylistic code changes
- slightly change logging/exceptions
- improved comments
- pin Omegaconf version
- remove faulty code

* remove default None args

* mark stuff private in `config`

* remove __main__ from `trainer_base`

* rework `load_streams` after comments

* update uv.lock, fix missing adjustment in evaluate

* cosmetic fixes

Author: grassesi

config/default_config.yml

@@ -0,0 +1,113 @@
+streams_directory: "./config/streams/streams_anemoi/"
+
+embed_orientation: "channels"
+embed_local_coords: True
+embed_centroids_local_coords: False
+embed_size_centroids: 64
+embed_unembed_mode: "block"
+
+target_cell_local_prediction: True
+target_coords_local: True
+
+ae_local_dim_embed: 1024
+ae_local_num_blocks: 2
+ae_local_num_heads: 16
+ae_local_dropout_rate: 0.1
+ae_local_with_qk_lnorm: True
+
+ae_local_num_queries: 2
+ae_local_queries_per_cell: False
+ae_adapter_num_heads: 16
+ae_adapter_embed: 128
+ae_adapter_with_qk_lnorm: True
+ae_adapter_with_residual: True
+ae_adapter_dropout_rate: 0.1
+
+ae_global_dim_embed: 2048
+ae_global_num_blocks: 8
+ae_global_num_heads: 32
+ae_global_dropout_rate: 0.1
+ae_global_with_qk_lnorm: True
+ae_global_att_dense_rate: 0.2
+ae_global_block_factor: 64
+ae_global_mlp_hidden_factor: 2
+
+pred_adapter_kv: False
+pred_self_attention: True
+pred_dyadic_dims: False
+pred_mlp_adaln: True
+
+forecast_delta_hrs: 0
+forecast_steps: 0
+forecast_policy: null
+forecast_freeze_model: False
+forecast_att_dense_rate: 0.25
+fe_num_blocks: 0
+fe_num_heads: 16
+fe_dropout_rate: 0.1
+fe_with_qk_lnorm: True
+
+healpix_level: 5
+
+with_mixed_precision: True
+with_flash_attention: True
+compile_model: False
+with_fsdp: True
+
+loss_fcts:
+  -
+    - "mse"
+    - 1.0
+loss_fcts_val:
+  -
+    - "mse"
+    - 1.0
+
+batch_size: 1
+batch_size_validation: 1
+
+masking_mode: "forecast"
+masking_rate: 0.0
+masking_rate_sampling: True
+sampling_rate_target: 1.0
+
+num_epochs: 24
+samples_per_epoch: 4096
+samples_per_validation: 512
+shuffle: True
+
+lr_scaling_policy: "sqrt"
+lr_start: 0.000001
+lr_max: 0.00003
+lr_final_decay: 0.000001
+lr_final: 0.0
+lr_steps_warmup: 256
+lr_steps_cooldown: 4096
+lr_policy_warmup: "cosine"
+lr_policy_decay: "linear"
+lr_policy_cooldown: "linear"
+
+grad_clip: 5.0
+weight_decay: 0.1
+norm_type: "LayerNorm"
+nn_module: "te"
+
+start_date: 201301010000
+end_date: 202012310000
+start_date_val: 202101010000
+end_date_val: 202201010000
+len_hrs: 6
+step_hrs: 6
+input_window_steps: 1
+
+val_initial: False
+
+loader_num_workers: 8
+log_validation: 0
+
+istep: 0
+run_history: []
+
+desc: ""
+data_loader_rng_seed: ???
+run_id: ???

pyproject.toml

@@ -24,6 +24,7 @@ dependencies = [
  'psutil',
  "flash-attn; sys_platform == 'linux'",
  "polars~=1.25.2",
+ "omegaconf~=2.3.0",
 ]
 
 [project.urls]

src/weathergen/__init__.py

@@ -12,15 +12,15 @@
 import sys
 import time
 import traceback
+from pathlib import Path
 
 import pandas as pd
 
+import weathergen.utils.config as config
 from weathergen.train.trainer import Trainer
-from weathergen.utils.config import Config, load_overwrite_conf, load_private_conf
 from weathergen.utils.logger import init_loggers
 
 
-####################################################################################################
 def evaluate():
     """
     Evaluation function for WeatherGenerator model.
@@ -79,26 +79,17 @@ def evaluate():
     )
     parser.add_argument(
         "--private_config",
-        type=str,
+        type=Path,
         default=None,
         help="Path to private configuration file for paths.",
     )
 
     args = parser.parse_args()
 
-    # get the paths from the private config
-    private_cf = load_private_conf(args.private_config)
-
     # TODO: move somewhere else
     init_loggers()
 
-    # load config: if run_id is full path, it loads from there
-    model_path = private_cf["model_path"] if "model_path" in private_cf.keys() else "./models"
-    cf = Config.load(args.run_id, args.epoch, model_path)
-
-    # add parameters from private (paths) config
-    for k, v in private_cf.items():
-        setattr(cf, k, v)
+    cf = config.load_config(args.private_config, args.run_id, args.epoch, None)
 
     cf.run_history += [(cf.run_id, cf.istep)]
 
@@ -154,7 +145,7 @@ def train_continue() -> None:
     )
     parser.add_argument(
         "--private_config",
-        type=str,
+        type=Path,
         default=None,
         help="Path to private configuration file for paths.",
     )
@@ -165,16 +156,10 @@ def train_continue() -> None:
     )
 
     args = parser.parse_args()
-    # get the paths from the private config
-    private_cf = load_private_conf(args.private_config)
 
-    # load config if specified
-    model_path = private_cf["model_path"] if "model_path" in private_cf.keys() else "./models"
-    cf = Config.load(args.run_id, args.epoch, model_path)
+    cf = config.load_config(args.private_config, args.run_id, args.epoch, None)
 
     # track history of run to ensure traceability of results
-    if "run_history" not in cf.__dict__:
-        cf.run_history = []
     cf.run_history += [(cf.run_id, cf.istep)]
 
     #########################
@@ -210,7 +195,7 @@ def train_continue() -> None:
         cf.istep = 0
 
     trainer = Trainer()
-    trainer.run(cf, private_cf, args.run_id, args.epoch, args.run_id_new)
+    trainer.run(cf, args.run_id, args.epoch, args.run_id_new)
 
 
 ####################################################################################################
@@ -235,171 +220,32 @@ def train() -> None:
     )
     parser.add_argument(
         "--private_config",
-        type=str,
+        type=Path,
         default=None,
         help="Path to private configuration file for paths",
     )
     parser.add_argument(
         "--config",
-        type=str,
+        type=Path,
         default=None,
-        help="Path to private configuration file for overwriting the defaults in the function body. Defaults to None.",
+        help="Optional experiment specfic configuration file",
     )
 
     args = parser.parse_args()
 
     # TODO: move somewhere else
     init_loggers()
 
-    # get the non-default configs: private and overwrite
-    private_cf = load_private_conf(args.private_config)
-    overwrite_cf = load_overwrite_conf(args.config)
-
-    cf = Config()
-
-    # directory where input streams are specified
-    # cf.streams_directory = './streams_large/'
-    cf.streams_directory = "./config/streams/streams_anemoi/"
-    # cf.streams_directory = "./config/streams/streams_mixed/"
-    # cf.streams_directory = "./streams_mixed/"
-
-    # embed_orientation : 'channels' or 'columns'
-    # channels: embedding is per channel for a token (#tokens=num_channels)
-    # columns:  embedding is per "column", all channels are embedded together (#tokens=token_size)
-    # the per-stream embedding paramters, in particular dim_embed, have to be chosen accordingly
-    cf.embed_orientation = "channels"
-    cf.embed_local_coords = True
-    # False since per cell coords are meaningless for cells
-    cf.embed_centroids_local_coords = False
-    cf.embed_size_centroids = 64
-    cf.embed_unembed_mode = "block"
-
-    cf.target_cell_local_prediction = True
-    cf.target_coords_local = True
-
-    # parameters for local assimilation engine
-    cf.ae_local_dim_embed = 1024  # 2048 #1024
-    cf.ae_local_num_blocks = 2
-    cf.ae_local_num_heads = 16
-    cf.ae_local_dropout_rate = 0.1
-    cf.ae_local_with_qk_lnorm = True
-
-    # assimilation engine local -> global adapter
-    cf.ae_local_num_queries = 2
-    cf.ae_local_queries_per_cell = False
-    cf.ae_adapter_num_heads = 16
-    cf.ae_adapter_embed = 128
-    cf.ae_adapter_with_qk_lnorm = True
-    cf.ae_adapter_with_residual = True
-    cf.ae_adapter_dropout_rate = 0.1
-
-    # parameters for global assimilation engine
-    cf.ae_global_dim_embed = 2048
-    cf.ae_global_num_blocks = 8
-    cf.ae_global_num_heads = 32
-    cf.ae_global_dropout_rate = 0.1
-    cf.ae_global_with_qk_lnorm = True
-    cf.ae_global_att_dense_rate = 0.2  # 0.25 : every 4-th block is dense attention
-    cf.ae_global_block_factor = 64
-    cf.ae_global_mlp_hidden_factor = 2
-
-    cf.pred_adapter_kv = False
-    cf.pred_self_attention = True
-    cf.pred_dyadic_dims = False
-    cf.pred_mlp_adaln = True
-
-    # forecasting engine
-    cf.forecast_delta_hrs = 0
-    cf.forecast_steps = 0  # [j for j in range(1,11) for i in range(1)]
-    cf.forecast_policy = None  #'fixed', 'sequential'
-    cf.forecast_freeze_model = False  # False
-    cf.forecast_att_dense_rate = 0.25
-
-    cf.fe_num_blocks = 0
-    cf.fe_num_heads = 16
-    cf.fe_dropout_rate = 0.1
-    cf.fe_with_qk_lnorm = True
-
-    cf.healpix_level = 5
-
-    # working precision
-    cf.with_mixed_precision = True
-    cf.with_flash_attention = True
+    cf = config.load_config(args.private_config, None, None, args.config)
+
     if cf.with_flash_attention:
         assert cf.with_mixed_precision
-    # compile entire model
-    cf.compile_model = False
-
-    cf.with_fsdp = True
-
-    cf.loss_fcts = [["mse", 1.0]]
-    cf.loss_fcts_val = [["mse", 1.0]]
-    # cf.loss_fcts = [['mse', 0.5], ['stats', 0.5]]
-    # cf.loss_fcts_val = [['mse', 0.5], ['stats', 0.5]]
-
-    cf.batch_size = 1
-    cf.batch_size_validation = 1
-
-    # forecast
-    cf.masking_mode = "forecast"
-    cf.masking_rate = 0.0
-    cf.masking_rate_sampling = True  # False
-    cf.sampling_rate_target = 1.0
-
-    cf.num_epochs = 24
-    cf.samples_per_epoch = 4096
-    cf.samples_per_validation = 512
-    cf.shuffle = True
-
-    cf.lr_scaling_policy = "sqrt"
-    cf.lr_start = 0.000001
-    cf.lr_max = 0.00003
-    cf.lr_final_decay = 0.000001
-    cf.lr_final = 0.0
-    cf.lr_steps_warmup = 256
-    cf.lr_steps_cooldown = 4096
-    cf.lr_policy_warmup = "cosine"
-    cf.lr_policy_decay = "linear"
-    cf.lr_policy_cooldown = "linear"
-
-    cf.grad_clip = 5.0
-    cf.weight_decay = 0.1
-    cf.norm_type = "LayerNorm"  #'LayerNorm' #'RMSNorm'
-    cf.nn_module = "te"
-
-    cf.start_date = 201301010000
-    cf.end_date = 202012310000
-    cf.start_date_val = 202101010000
-    cf.end_date_val = 202201010000
-    cf.len_hrs = 6
-    cf.step_hrs = 6
-    cf.input_window_steps = 1
-
-    cf.val_initial = False
-
-    cf.loader_num_workers = 8
     cf.data_loader_rng_seed = int(time.time())
-    cf.log_validation = 0
-
-    cf.istep = 0
-    cf.run_history = []
-
-    cf.run_id = args.run_id
-    cf.desc = ""
-
-    # overwrite parameters from private config
-    for k, v in private_cf.items():
-        setattr(cf, k, v)
-    cf.data_path = private_cf["data_path_anemoi"]  # for backward compatibility
-
-    # overwrite parameters from overwrite config
-    for k, v in overwrite_cf.items():
-        setattr(cf, k, v)
 
     trainer = Trainer(log_freq=20, checkpoint_freq=250, print_freq=10)
 
     try:
-        trainer.run(cf, private_cf)
+        trainer.run(cf)
     except Exception:
         extype, value, tb = sys.exc_info()
         traceback.print_exc()