From 7fab8e9361c9007810aacfeb75faefebb9f9d7ca Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Wed, 18 Feb 2026 13:24:46 -0500 Subject: [PATCH 1/2] case informaiton --- commands.txt | 6 ++ config/training_config_experiment_2.json | 131 +++++++++++++++++++++++ docs/tropical_bad_vars_observations.md | 70 ++++++++++++ 3 files changed, 207 insertions(+) create mode 100644 config/training_config_experiment_2.json create mode 100644 docs/tropical_bad_vars_observations.md diff --git a/commands.txt b/commands.txt index 377acd0..ce13eb2 100644 --- a/commands.txt +++ b/commands.txt @@ -160,4 +160,10 @@ python train_cnp_model.py \ --variable-list CNP_IO_updated9_dev_dw.txt \ --training-config config/training_config_experiment_2_plus.json \ --epoch 80 \ + --tropical-only + + + python train_cnp_model.py \ + --variable-list CNP_IO_updated9_dev_dw.txt \ + --training-config config/training_config_experiment_2.json \ --tropical-only \ No newline at end of file diff --git a/config/training_config_experiment_2.json b/config/training_config_experiment_2.json new file mode 100644 index 0000000..7a98a61 --- /dev/null +++ b/config/training_config_experiment_2.json @@ -0,0 +1,131 @@ +{ + "training_hyperparameters": { + "num_epochs": 100, + "batch_size": 128, + "learning_rate": 0.0001, + "optimizer_type": "adam", + "weight_decay": 0.0, + "use_scheduler": false, + "xsmrpool_loss_weight": 10.0, + "litter_c_loss_weight": 1.0, + "litter_n_loss_weight": 1.0, + "litter_p_loss_weight": 1.0, + "scalar_loss_weight": 1.0, + "vector_loss_weight": 1.0, + "matrix_loss_weight": 1.0 + }, + "reproducibility_config": { + "random_seed": 42, + "strict_determinism": false, + "train_split": 0.8, + "normalization": "individual", + "dropout_p": null + }, + "data_filtering_config": { + "tropical_only": true, + "tropical_lat_range": [-23.5, 23.5] + }, + "variable_weights": { + "pft1d_weights": { + "cpool": 6, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemc": 6, + "deadstemn": 6, + "deadstemp": 5, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "frootc_storage": 3, + "frootn_storage": 3, + "frootp_storage": 3, + "tlai": 4 + }, + "soil2d_weights": { + "primp_vr": 9, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7, + "litr3n_vr": 5, + "labilep_vr": 5, + "secondp_vr": 2.5, + "solutinorganicp_vr": 5, + "litr3p_vr": 5 + }, + "scalar_weights": { + "GPP": 1.5, + "NPP": 1.5, + "AR": 1.2, + "HR": 1.2 + } + }, + "tail_aware_config": { + "loss": "log1p_mse", + "epsilon": 1e-8 + }, + "tail_aware_weights": { + "cpool": 6, + "deadstemc": 6, + "deadcrootc": 5, + "livestemc": 5, + "livecrootc": 5, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemp": 5, + "deadstemn": 6, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "labilep": 5, + "labilep_vr": 5, + "primp_vr": 9, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "litr3c_vr": 5, + "litr3p_vr": 5, + "litr3n_vr": 5, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7 + }, + "pft_mask_config": { + "mask_absent_pfts": true, + "pft_presence_threshold": 0.0 + }, + "pft_zero_sparsity_config": { + "weight": 1.0, + "threshold": 1e-8 + }, + "pft_zero_sparsity_weights": { + "cpool": 2, + "deadstemc": 1, + "deadcrootc": 1, + "livestemc": 1, + "livecrootc": 1 + }, + "pft1d_activation_overrides": { + "cpool": "abs", + "deadstemc": "abs", + "deadcrootc": "abs", + "livestemc": "abs", + "livecrootc": "abs" + } +} diff --git a/docs/tropical_bad_vars_observations.md b/docs/tropical_bad_vars_observations.md new file mode 100644 index 0000000..63a2816 --- /dev/null +++ b/docs/tropical_bad_vars_observations.md @@ -0,0 +1,70 @@ +# Tropical Bad-Variable Observations and General Suggestions + +## Scope +- Dataset: tropical-only subset (Latitude between -23.5 and 23.5). +- Source: Trendy_1 training data (21 PKL files). +- Samples: 4,030 tropical rows. +- Variables analyzed: `npool`, `ppool`, `cpool`, `totvegc`, `leafc/n/p`, `deadstemc/n`, `primp_vr`, `litr2c/n/p_vr`, `soil1c/n/p_vr`. + +## Tropical Distribution Summary (aggregate across PFTs/layers) +Values below are computed over all PFTs (for 1D PFT variables) or layers (for 2D soil variables). + +- `Y_npool`: nonzero 46.38%, max 151, median(nz) 10, q90 38.7, q95 53.4 +- `Y_ppool`: nonzero 46.37%, max 8.16, median(nz) 1, q90 1.9, q95 3.24 +- `Y_cpool`: nonzero 9.83%, max 8.57e3, median(nz) 748, q90 2.59e3, q95 3.87e3 +- `Y_totvegc`: nonzero 45.83%, max 6.95e4, median(nz) 1.1, q90 4.48e3, q95 1.83e4 +- `Y_leafc`: nonzero 21.47%, max 811, median(nz) 1, q90 335, q95 420 +- `Y_leafn`: nonzero 21.47%, max 32.4, median(nz) 0.0333, q90 12.1, q95 16.2 +- `Y_leafp`: nonzero 21.47%, max 2.16, median(nz) 0.0025, q90 0.816, q95 1.17 +- `Y_deadstemc`: nonzero 31.92%, max 5.13e4, median(nz) 1, q90 3.79e3, q95 1.61e4 +- `Y_deadstemn`: nonzero 31.92%, max 103, median(nz) 0.002, q90 7.59, q95 32.2 +- `Y_primp_vr`: nonzero 66.67%, max 1.88e3, median(nz) 184, q90 1.26e3, q95 1.26e3 +- `Y_litr2c_vr`: nonzero 3.87%, max 7.98e3, median(nz) 27.6, q90 478, q95 807 +- `Y_litr2n_vr`: nonzero 3.87%, max 167, median(nz) 0.48, q90 6.26, q95 13.2 +- `Y_litr2p_vr`: nonzero 3.96%, max 9.6, median(nz) 0.0271, q90 0.511, q95 0.888 +- `Y_soil1c_vr`: nonzero 3.74%, max 779, median(nz) 6.86, q90 72.2, q95 108 +- `Y_soil1n_vr`: nonzero 3.74%, max 64.9, median(nz) 0.572, q90 6.01, q95 8.97 +- `Y_soil1p_vr`: nonzero 4.01%, max 2.16, median(nz) 0.0163, q90 0.193, q95 0.288 + +## Why These Variables Perform Poorly (Patterns) +1. **High sparsity**: soil1* and litr2* are ~4% non‑zero, so the model mostly sees zeros. +2. **Heavy‑tail distributions**: cpool, deadstemc, totvegc have huge max vs median, so typical losses underpredict extremes. +3. **Near‑constant regimes**: npool ~10 and ppool ~1 for many PFTs; the model learns the constant and misses tails. +4. **Mixed scales**: leafp/leafn are tiny vs leafc/cpool; shared heads can bias toward larger scales. + +## General Suggestions (Applies to All These Variables) + +### A) Target tail behavior only where it exists +- Apply tail‑aware loss **only** to tail‑heavy variables: + - Strong tails: `cpool`, `deadstemc`, `totvegc`, `primp_vr`, `litr2c_vr`, `soil1c_vr` + - Moderate tails: `npool`, `ppool`, `deadstemn`, `leafc` +- Use `log1p_huber` or `log1p_quantile` to push higher values without exploding gradients. + +### B) Separate sparse vs dense handling +- For very sparse 2D soil variables (`litr2*`, `soil1*`), use **zero‑inflated loss**: + - Loss = BCE/hinge for zero vs non‑zero + regression loss for non‑zero. +- Alternatively, add a **non‑zero mask head** and only regress where mask=1. + +### C) Per‑variable or per‑group weighting, not global +- Avoid global weight increases (past runs degraded overall quality). +- Increase weights only for a **short list of tail variables** and keep others at baseline. + +### D) Sampling strategy +- Oversample tropical rows where target values are in the top 10% (tail enrichment). +- This is less disruptive than large loss weights. + +### E) Scaling / normalization improvements +- For variables with extreme max/median ratios (cpool, deadstemc, totvegc), consider: + - log1p transform before scaling + - per‑variable robust scaling (median/IQR) + +### F) Optional architecture tweaks +- Add **per‑variable heads** for the most problematic groups: + - One head for soil1/litr2 (sparse 2D) + - One head for pool variables (cpool/npool/ppool) + +## Concrete Next Step (Minimal Risk) +- Keep experiment2 as baseline. +- Add tail‑aware loss **only** for `cpool`, `deadstemc`, `totvegc`, `primp_vr`, `litr2c_vr`, `soil1c_vr`. +- Keep `npool` and `ppool` weights unchanged for now, but add tail‑aware for their known tail PFTs (see `npool_ppool_tropical_observations.md`). + From c5b19246d8020ab98379458a7781cc213a1075bb Mon Sep 17 00:00:00 2001 From: Dali Wang Date: Wed, 18 Feb 2026 16:41:06 -0500 Subject: [PATCH 2/2] reference experiment2 repeat --- CNP_IO_updated9_dev_dw_repeat.txt | 75 ++++++++++ commands.txt | 6 +- .../training_config_experiment_2_repeat.json | 135 ++++++++++++++++++ train_cnp_model.py | 51 +++---- 4 files changed, 240 insertions(+), 27 deletions(-) create mode 100644 CNP_IO_updated9_dev_dw_repeat.txt create mode 100644 config/training_config_experiment_2_repeat.json diff --git a/CNP_IO_updated9_dev_dw_repeat.txt b/CNP_IO_updated9_dev_dw_repeat.txt new file mode 100644 index 0000000..ef3f5b1 --- /dev/null +++ b/CNP_IO_updated9_dev_dw_repeat.txt @@ -0,0 +1,75 @@ +# Dataset roots (any absolute paths) +TRENDY1_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_1_data_CNP +TRENDY05_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData/Trendy_05_data_CNP +TVA4KM_PATH: /mnt/proj-shared/AI4BGC_7xw/TrainingData/TVA_4km_data_CNP + +# Optional extra roots (comma-separated) +DATA_PATHS: /another/path1, /another/path2 + +# Global fallback pattern if a dataset-specific one isn't set +#FILE_PATTERN: enhanced_1_training_data_batch_*.pkl +FILE_PATTERN: 1_training_data_batch_*.pkl + +# Per-dataset patterns (overrides FILE_PATTERN for that path only) +TVA4KM_FILE_PATTERN: enhanced_monthly_training_data_batch_*.pkl + + +LONGITUDE FILTERING - 2 longitudes: +• 0, 358.75 + +TIME SERIES VARIABLES (Climate Forcing) - 6 variables: +• FLDS, PSRF, FSDS, QBOT, PRECTmms, TBOT + +SURFACE PROPERTIES - 49 variables: +• Latitude, Longitude, AREA, landfrac, LANDFRAC_PFT, PCT_NATVEG + +• OCCLUDED_P, SECONDARY_P, LABILE_P, APATITE_P + +• SOIL_COLOR, SOIL_ORDER + +• PCT_NAT_PFT_0, PCT_NAT_PFT_1, PCT_NAT_PFT_2, PCT_NAT_PFT_3, PCT_NAT_PFT_4, PCT_NAT_PFT_5, PCT_NAT_PFT_6, PCT_NAT_PFT_7, PCT_NAT_PFT_8 +• PCT_NAT_PFT_9, PCT_NAT_PFT_10, PCT_NAT_PFT_11, PCT_NAT_PFT_12, PCT_NAT_PFT_13, PCT_NAT_PFT_14, PCT_NAT_PFT_15, PCT_NAT_PFT_16 + +• PCT_CLAY_0, PCT_CLAY_1, PCT_CLAY_2, PCT_CLAY_3, PCT_CLAY_4, PCT_CLAY_5, PCT_CLAY_6, PCT_CLAY_7, PCT_CLAY_8, PCT_CLAY_9 +• PCT_SAND_0, PCT_SAND_1, PCT_SAND_2, PCT_SAND_3, PCT_SAND_4, PCT_SAND_5, PCT_SAND_6, PCT_SAND_7, PCT_SAND_8, PCT_SAND_9 + +PFT PARAMETERS (Plant Functional Type Characteristics) - 44 variables: + +• pft_deadwdcn, pft_frootcn, pft_leafcn, pft_lflitcn, pft_livewdcn, pft_c3psn, pft_croot_stem, pft_crop, pft_dleaf +• pft_dsladlai, pft_evergreen, pft_fcur, pft_flivewd, pft_flnr, pft_fr_fcel, pft_fr_flab, pft_fr_flig, pft_froot_leaf +• pft_grperc, pft_grpnow, pft_leaf_long, pft_lf_fcel, pft_lf_flab, pft_lf_flig, pft_rholnir, pft_rholvis, pft_rhosnir, pft_rhosvis +• pft_roota_par, pft_rootb_par, pft_rootprof_beta, pft_season_decid, pft_slatop, pft_smpsc, pft_smpso, pft_stem_leaf, pft_stress_decid +• pft_taulnir, pft_taulvis, pft_tausnir, pft_tausvis, pft_woody, pft_xl, pft_z0mr + +SCALAR VARIABLES (1D - 4 variables): +• GPP, NPP, AR, HR + +1D PFT VARIABLES (41 variables): + +• deadcrootc, deadcrootc_storage, deadcrootn, deadcrootn_storage, deadcrootp, deadcrootp_storage +• deadstemc, deadstemc_storage, deadstemn, deadstemn_storage, deadstemp, deadstemp_storage + +• leafc, leafc_storage, leafn, leafn_storage, leafp, leafp_storage +• frootc, frootc_storage, frootn, frootn_storage, frootp, frootp_storage + +• livestemc, livestemc_storage, livestemn, livestemn_storage, livestemp, livestemp_storage, +• livecrootc, livecrootc_storage, livecrootn, livecrootn_storage, livecrootp, livecrootp_storage + +• cpool, npool, ppool + +• tlai, totvegc + +2D VARIABLES (layered - 25 variables): + +• cwdc_vr, cwdn_vr, cwdp_vr + +• litr2c_vr, litr3c_vr +• litr2n_vr, litr3n_vr +• litr2p_vr, litr3p_vr + +• soil1c_vr, soil1n_vr, soil1p_vr +• soil2c_vr, soil2n_vr, soil2p_vr +• soil3c_vr, soil3n_vr, soil3p_vr +• soil4c_vr, soil4n_vr, soil4p_vr + +• labilep_vr , occlp_vr, primp_vr, secondp_vr, solutionp_vr diff --git a/commands.txt b/commands.txt index ce13eb2..ba2ae04 100644 --- a/commands.txt +++ b/commands.txt @@ -112,7 +112,7 @@ python train_cnp_model.py \ --variable-list CNP_IO_updated9_dev_dw.txt \ --training-config-json config/training_config_experiment_2.json \ --pft-zero-sparsity-weight 1.0 \ - --epoch 50 \ + --epoch 80 \ --tropical-only \ 2>&1 | tee training_cpool_experiment.log & @@ -164,6 +164,6 @@ python train_cnp_model.py \ python train_cnp_model.py \ - --variable-list CNP_IO_updated9_dev_dw.txt \ - --training-config config/training_config_experiment_2.json \ + --variable-list CNP_IO_updated9_dev_dw_repeat.txt \ + --training-config config/training_config_experiment_2_repeat.json \ --tropical-only \ No newline at end of file diff --git a/config/training_config_experiment_2_repeat.json b/config/training_config_experiment_2_repeat.json new file mode 100644 index 0000000..876c295 --- /dev/null +++ b/config/training_config_experiment_2_repeat.json @@ -0,0 +1,135 @@ +{ + "_description": "Repeat of run_20260212_162802_experiment_2: 100 epochs, AdamW, cosine scheduler. Use this config (with same variable list) to reproduce experiment_2 behavior.", + "training_hyperparameters": { + "num_epochs": 100, + "batch_size": 128, + "learning_rate": 0.0001, + "optimizer_type": "adamw", + "weight_decay": 0.01, + "use_scheduler": true, + "scheduler_type": "cosine", + "scheduler_step_size": 10, + "scheduler_gamma": 0.1, + "xsmrpool_loss_weight": 10, + "litter_c_loss_weight": 1, + "litter_n_loss_weight": 1, + "litter_p_loss_weight": 1, + "scalar_loss_weight": 1, + "vector_loss_weight": 1, + "matrix_loss_weight": 1 + }, + "reproducibility_config": { + "random_seed": 42, + "strict_determinism": false, + "train_split": 0.8, + "normalization": "individual", + "dropout_p": null + }, + "data_filtering_config": { + "tropical_only": true, + "tropical_lat_range": [-23.5, 23.5] + }, + "variable_weights": { + "pft1d_weights": { + "cpool": 6, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemc": 6, + "deadstemn": 6, + "deadstemp": 5, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "frootc_storage": 3, + "frootn_storage": 3, + "frootp_storage": 3, + "tlai": 4 + }, + "soil2d_weights": { + "primp_vr": 9, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7, + "litr3n_vr": 5, + "labilep_vr": 5, + "secondp_vr": 2.5, + "solutinorganicp_vr": 5, + "litr3p_vr": 5 + }, + "scalar_weights": { + "GPP": 1.5, + "NPP": 1.5, + "AR": 1.2, + "HR": 1.2 + } + }, + "tail_aware_config": { + "loss": "log1p_mse", + "epsilon": 1e-8 + }, + "tail_aware_weights": { + "cpool": 6, + "deadstemc": 6, + "deadcrootc": 5, + "livestemc": 5, + "livecrootc": 5, + "npool": 8, + "ppool": 8, + "totvegc": 6, + "leafc": 6, + "leafn": 6, + "leafp": 6, + "leafc_storage": 5, + "leafn_storage": 5, + "leafp_storage": 5, + "deadstemp": 5, + "deadstemn": 6, + "frootc": 4, + "frootn": 4, + "frootp": 4, + "labilep": 5, + "labilep_vr": 5, + "primp_vr": 9, + "litr2c_vr": 8, + "litr2p_vr": 8, + "litr2n_vr": 8, + "litr3c_vr": 5, + "litr3p_vr": 5, + "litr3n_vr": 5, + "soil1c_vr": 7, + "soil1n_vr": 7, + "soil1p_vr": 7 + }, + "pft_mask_config": { + "mask_absent_pfts": true, + "pft_presence_threshold": 0 + }, + "pft_zero_sparsity_config": { + "weight": 1, + "threshold": 1e-8 + }, + "pft_zero_sparsity_weights": { + "cpool": 2, + "deadstemc": 1, + "deadcrootc": 1, + "livestemc": 1, + "livecrootc": 1 + }, + "pft1d_activation_overrides": { + "cpool": "abs", + "deadstemc": "abs", + "deadcrootc": "abs", + "livestemc": "abs", + "livecrootc": "abs" + } +} diff --git a/train_cnp_model.py b/train_cnp_model.py index c7116a9..f1a4e80 100644 --- a/train_cnp_model.py +++ b/train_cnp_model.py @@ -57,14 +57,16 @@ def setup_logging(log_file: str, level: str = 'INFO') -> None: """Set up logging configuration with a specific log file.""" - logging.basicConfig( - level=getattr(logging, level.upper()), - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', - handlers=[ - logging.StreamHandler(sys.stdout), - logging.FileHandler(log_file) - ] - ) + level_value = getattr(logging, level.upper(), logging.INFO) + fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s' + handlers = [ + logging.StreamHandler(sys.stdout), + logging.FileHandler(log_file) + ] + try: + logging.basicConfig(level=level_value, format=fmt, handlers=handlers, force=True) + except TypeError: + logging.basicConfig(level=level_value, format=fmt, handlers=handlers) def set_global_determinism(seed: int) -> None: @@ -135,8 +137,8 @@ def main(): '--epochs', '--epoch', dest='epochs', type=int, - default=150, - help='Number of training epochs' + default=None, + help='Number of training epochs (default: use value from --training-config-json, or 50)' ) parser.add_argument( '--batch-size', @@ -406,9 +408,9 @@ def main(): output_dir = Path(args.output_dir) / f"run_{timestamp}" output_dir.mkdir(parents=True, exist_ok=True) - # Setup logging with timestamped log file in output directory + # Setup logging with timestamped log file in output directory (default INFO so log file is populated) log_file = output_dir / f"cnp_training_{timestamp}.log" - setup_logging(str(log_file), args.log_level if args.log_level else 'WARNING') + setup_logging(str(log_file), args.log_level or 'INFO') logger = logging.getLogger(__name__) logger.info(f"Output directory: {output_dir}") logger.info(f"Normalization method: {args.normalization}") @@ -655,20 +657,20 @@ def main(): if 'scheduler_gamma' in hyperparams: update_kwargs['scheduler_gamma'] = float(hyperparams['scheduler_gamma']) - # Loss weights (CLI takes precedence) - if 'scalar_loss_weight' in hyperparams and args.scalar_loss_weight is None: + # Loss weights (CLI takes precedence; use getattr so missing CLI args don't crash) + if 'scalar_loss_weight' in hyperparams and getattr(args, 'scalar_loss_weight', None) is None: update_kwargs['scalar_loss_weight'] = float(hyperparams['scalar_loss_weight']) - if 'vector_loss_weight' in hyperparams and args.vector_loss_weight is None: + if 'vector_loss_weight' in hyperparams and getattr(args, 'vector_loss_weight', None) is None: update_kwargs['vector_loss_weight'] = float(hyperparams['vector_loss_weight']) - if 'matrix_loss_weight' in hyperparams and args.matrix_loss_weight is None: + if 'matrix_loss_weight' in hyperparams and getattr(args, 'matrix_loss_weight', None) is None: update_kwargs['matrix_loss_weight'] = float(hyperparams['matrix_loss_weight']) - if 'xsmrpool_loss_weight' in hyperparams and args.xsmrpool_loss_weight is None: + if 'xsmrpool_loss_weight' in hyperparams and getattr(args, 'xsmrpool_loss_weight', None) is None: update_kwargs['xsmrpool_loss_weight'] = float(hyperparams['xsmrpool_loss_weight']) - if 'litter_c_loss_weight' in hyperparams and args.litter_c_loss_weight is None: + if 'litter_c_loss_weight' in hyperparams and getattr(args, 'litter_c_loss_weight', None) is None: update_kwargs['litter_c_loss_weight'] = float(hyperparams['litter_c_loss_weight']) - if 'litter_n_loss_weight' in hyperparams and args.litter_n_loss_weight is None: + if 'litter_n_loss_weight' in hyperparams and getattr(args, 'litter_n_loss_weight', None) is None: update_kwargs['litter_n_loss_weight'] = float(hyperparams['litter_n_loss_weight']) - if 'litter_p_loss_weight' in hyperparams and args.litter_p_loss_weight is None: + if 'litter_p_loss_weight' in hyperparams and getattr(args, 'litter_p_loss_weight', None) is None: update_kwargs['litter_p_loss_weight'] = float(hyperparams['litter_p_loss_weight']) if update_kwargs: @@ -707,8 +709,8 @@ def main(): # We'll apply this after checking if normalization was explicitly set via CLI repro_config['_normalization_from_config'] = norm_val - # Dropout (CLI takes precedence) - if 'dropout_p' in repro_config and args.dropout_p is None: + # Dropout (CLI takes precedence; skip if null in JSON) + if 'dropout_p' in repro_config and args.dropout_p is None and repro_config['dropout_p'] is not None: update_model_kwargs['dropout_p'] = float(repro_config['dropout_p']) if update_training_kwargs: @@ -842,10 +844,11 @@ def main(): logger.info(f"Using MAX_FILES from CNP_IO file: {config.data_config.max_files}") # Turn off GPU monitoring and debug logging config.update_training_config(log_gpu_memory=False, log_gpu_utilization=False) - # Override training parameters if specified + # Override training parameters if specified (CLI overrides config; if CLI not set, use config) effective_lr = args.learning_rate if args.learning_rate is not None else config.training_config.learning_rate + effective_epochs = args.epochs if args.epochs is not None else config.training_config.num_epochs config.update_training_config( - num_epochs=args.epochs, + num_epochs=effective_epochs, batch_size=args.batch_size, learning_rate=effective_lr, model_save_path=str(output_dir / "cnp_model.pt"),