diff --git a/climanet/train.py b/climanet/train.py index 682133f..c5ed08d 100644 --- a/climanet/train.py +++ b/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 diff --git a/scripts/README.md b/scripts/README.md index 739642d..3ee2566 100644 --- a/scripts/README.md +++ b/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//eso4clima/train/sst_01/`. + diff --git a/scripts/training.py b/scripts/training.py new file mode 100644 index 0000000..dc4c98a --- /dev/null +++ b/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, + ) diff --git a/scripts/training.slurm b/scripts/training.slurm new file mode 100644 index 0000000..0e8d7ce --- /dev/null +++ b/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.************"