25 tarin on three year data

climanet/train.py

@@ -25,6 +25,9 @@ class TrainConfig:
     patience: int = 10
     accumulation_steps: int = 1
     optimizer_lr: float = 1e-3
+    optimizer_weight_decay: float = 1e-2
+    scheduler_lr_factor: float = 0.1
+    scheduler_min_lr: float = 1e-5
     device: str = "cpu"
     verbose: bool = False
     verbose_epoch_interval: int = 20
@@ -104,7 +107,9 @@ def train_monthly_model(
 
     # Set the optimizer
     optimizer = torch.optim.AdamW(
-        model.parameters(), lr=training_config.optimizer_lr, weight_decay=1e-2
+        model.parameters(),
+        lr=training_config.optimizer_lr,
+        weight_decay=training_config.optimizer_weight_decay,
     )
 
     best_loss = float("inf")
@@ -120,9 +125,9 @@ def train_monthly_model(
     scheduler = ReduceLROnPlateau(
         optimizer,
         mode="min",
-        factor=0.5,
+        factor=training_config.scheduler_lr_factor,
         patience=training_config.patience // 2,  # Reduce LR before early stop triggers
-        min_lr=1e-7,
+        min_lr=training_config.scheduler_min_lr,
     )
 
     model.train()
@@ -215,6 +220,8 @@ def train_monthly_model(
         if counter >= training_config.patience and current_lr <= scheduler.min_lrs[0]:
             if training_config.store_logs:
                 writer.add_text("Training", f"Early stop at epoch {epoch}", epoch)
+            if training_config.verbose:
+                print(f"Early stopping triggered at epoch {epoch}. Best loss: {best_loss:.6f}")
             break
 
     # Restore best model

scripts/README.md

@@ -5,6 +5,7 @@
 - `data_preparation.*`: Scripts for preparing the data for training, tuning and evaluation and saving them to Zarr storage with specific chunking strategies.
 - `tuning.*`: Scripts for hyperparameter tuning.
 - `run_best_tuned_model.*`: Scripts for running the best tuned model on the test set.
+- `training.*`: Scripts for training the model with the best hyperparameters found in the tuning experiments.
 
 ## Experiments
 
@@ -25,3 +26,17 @@
     'optimizer_lr': 0.001787422899066508,
     'batch_config': {'batch_size': 100, 'accumulation_steps': 2}}
  ```
+
+### Training experiments
+
+Use the best hyperparameters found in the tuning experiments to train the model
+on the training set; three years 2018-2020 for training, 2021 for validation in
+the training loop. Because three years of hourly data is too large to fit into
+memory, we use `load_lazy` option in the dataset. This makes the training
+process slower, but allows us to train on larger-than-memory datasets. The
+training is done for 100 epochs, and the best model is saved based on the
+validation loss. Note that the training is done only on one GPU node, including
+4 GPUs.
+
+The results are stored at `/work/<account_id>/eso4clima/train/sst_01/`.
+

scripts/training.py

@@ -0,0 +1,175 @@
+import argparse
+from pathlib import Path
+
+import ray
+import xarray as xr
+
+from climanet.dataset import DataLoaderConfig, STDataset
+from climanet.st_encoder_decoder import SpatioTemporalModel
+from climanet.train import TrainConfig, train_monthly_model
+from climanet.utils import configure_compute_resources, read_st_data, set_seed
+
+
+def _build_dataset(
+    prepared_data_dir: Path,
+    years: list[int],
+    var_name: str,
+    land_mask: xr.DataArray,
+    patch_size: tuple[int, int, int],
+    stride: tuple[int, int],
+) -> STDataset:
+
+    data = [read_st_data(data_path=f"{prepared_data_dir}/{year}", var_name=var_name) for year in years]
+    input_das, input_da_nan_masks, monthly_das, padded_days_masks, time_features_list = zip(*data)
+
+    input_da = xr.concat(input_das, dim="M")
+    input_da_nan_mask = xr.concat(input_da_nan_masks, dim="M")
+    monthly_da = xr.concat(monthly_das, dim="M")
+    padded_days_mask = xr.concat(padded_days_masks, dim="M")
+    time_features = xr.concat(time_features_list, dim="M")
+
+    return STDataset(
+        input_da=input_da,
+        input_da_nan_mask=input_da_nan_mask,
+        monthly_da=monthly_da,
+        padded_days_mask=padded_days_mask,
+        time_features=time_features,
+        land_mask=land_mask,
+        patch_size=patch_size,
+        stride=stride,
+        sh_embed_dim=96,
+        sh_order_L=10,
+        verbose=False,
+        load_lazy=True,
+    )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--run-dir",
+        type=str,
+        default=Path("./run_dir").resolve(),
+    )
+    parser.add_argument(
+        "--prepared-data-dir",
+        type=str,
+        default=Path("./data").resolve(),
+    )
+    parser.add_argument(
+        "--tune-dir",
+        type=str,
+        default=Path("./data").resolve(),
+    )
+    parser.add_argument(
+        "--lsm-dir",
+        type=str,
+        default=Path("./data").resolve(),
+    )
+    args = parser.parse_args()
+
+    var_name = "tos"
+    device = "cuda"
+    prepared_data_dir = Path(args.prepared_data_dir).resolve()
+    lsm_dir = Path(args.lsm_dir).resolve()
+    tune_dir = Path(args.tune_dir).resolve()
+    run_dir = Path(args.run_dir).resolve()
+
+    # Load the best hyperparameters from tuning
+    analysis = ray.tune.ExperimentAnalysis(str(tune_dir))
+    best_result = analysis.get_best_trial("loss", "min")
+    best_config = best_result.config
+
+    # set the random seed for reproducibility
+    set_seed()
+
+    # Build dataset for training and validation
+    lsm_file_path = lsm_dir / "era5_lsm_bool.nc"
+    lsm_mask = xr.open_dataset(lsm_file_path)["lsm"]  # make sure is dask array
+
+    dataset_patch_size = (1, 40, 40)
+    dataset_stride = (20, 20)
+
+    train_years = [2018, 2019, 2020]
+    dataset_train = _build_dataset(
+        prepared_data_dir=prepared_data_dir,
+        years=train_years,
+        var_name=var_name,
+        land_mask=lsm_mask,
+        patch_size=dataset_patch_size,
+        stride=dataset_stride,
+    )
+
+    validation_year = [2021]
+    dataset_validation = _build_dataset(
+        prepared_data_dir=prepared_data_dir,
+        years=validation_year,
+        var_name=var_name,
+        land_mask=lsm_mask,
+        patch_size=dataset_patch_size,
+        stride=dataset_stride,
+    )
+
+    # Build the dataloader config
+    dataloader_num_workers = 32  # adjust if needed
+    use_cuda = device == "cuda"
+    dataloader_config = DataLoaderConfig(
+        batch_size=100, # adjust if OOM issue
+        shuffle=True,
+        num_workers=dataloader_num_workers,
+        pin_memory=use_cuda,
+        persistent_workers=True,
+        device=device,
+        multiprocessing_context="spawn",
+    )
+
+    # Build the model with the best hyperparameters from tuning
+    patch_size = (1, best_config["patch_size"], best_config["patch_size"])
+    overlap = best_config["overlap"]
+    embed_dim = best_config["embed_dim"]
+    dropout = best_config["dropout"]
+    hidden = best_config["hidden"]
+    spatial_depth = best_config["spatial_depth"]
+    spatial_heads = best_config["spatial_heads"]
+
+    model = SpatioTemporalModel(
+        patch_size=patch_size,
+        overlap=overlap,
+        embed_dim=embed_dim,
+        dropout=dropout,
+        hidden=hidden,
+        spatial_depth=spatial_depth,
+        spatial_heads=spatial_heads,
+    )
+
+    # move the model to GPU and configure compute resources
+    model = configure_compute_resources(
+        model,
+        device=device,
+        compute_threads=None,  # on gpu, it is not used
+        dataloader_num_workers=dataloader_num_workers
+    )
+
+    # Training configuration
+    training_config = TrainConfig(
+        calculate_residuals=True,
+        num_epoch=101,
+        patience=10,
+        accumulation_steps=2,
+        optimizer_lr=best_config["optimizer_lr"],
+        device=device,
+        verbose=True,
+        verbose_epoch_interval=20,
+        tune_checkpoint=False,
+        store_model=True,
+        store_logs=True,
+    )
+
+    trained_model = train_monthly_model(
+        model=model,
+        dataset_train=dataset_train,
+        dataloader_config=dataloader_config,
+        training_config=training_config,
+        dataset_validation=dataset_validation,
+        run_dir=run_dir,
+    )

scripts/training.slurm

@@ -0,0 +1,37 @@
+#!/bin/bash
+#SBATCH --job-name=training
+#SBATCH --partition=gpu
+#SBATCH --constraint=a100_80
+#SBATCH --nodes=1
+#SBATCH --ntasks-per-node=1
+#SBATCH --cpus-per-task=128
+#SBATCH --gpus-per-task=4
+#SBATCH --exclusive
+#SBATCH --mem=0
+#SBATCH --time=12:00:00
+#SBATCH --account=bd0854
+#SBATCH --output=training_%j.out
+
+set -euo pipefail
+ulimit -s 204800
+
+# Activate uv env
+UV_ENV="$HOME/climanet_py314"
+source "$UV_ENV/bin/activate"
+
+# Set the scratch directory because they are avialable to all nodes
+RUN_DIR="/scratch/b/$USER/train"
+
+# data directory (adjust this path to your data location)
+PREPARED_DATA_DIR="/scratch/b/$USER/data"
+TUNE_DIR="/scratch/b/$USER/tune"
+LSM_DIR="/scratch/b/$USER/data"
+
+echo "Starting training.py script..."
+python -u $HOME/ClimaNet/scripts/training.py \
+  --run-dir "$RUN_DIR" \
+  --prepared-data-dir "$PREPARED_DATA_DIR" \
+  --tune-dir "$TUNE_DIR" \
+  --lsm-dir "$LSM_DIR"
+
+echo "**********Training script completed.************"