diff --git a/climanet/dataset.py b/climanet/dataset.py index 5cf5306..50571a8 100644 --- a/climanet/dataset.py +++ b/climanet/dataset.py @@ -110,8 +110,8 @@ def __init__( ) # (M, T=31, 3) # Store coordinate arrays - self.lat_coords = input_da[spatial_dims[0]].to_numpy().copy() - self.lon_coords = input_da[spatial_dims[1]].to_numpy().copy() + self.lat_coords = torch.from_numpy(input_da[spatial_dims[0]].to_numpy().copy()) + self.lon_coords = torch.from_numpy(input_da[spatial_dims[1]].to_numpy().copy()) if land_mask is not None: lm = torch.from_numpy(land_mask.values.copy()).bool() @@ -146,6 +146,7 @@ def __init__( self.patch_geo_embeddings, self.patch_scale_features = ( self._compute_geoscalepatch_embeddings() ) + self.scale_f_dim = torch.tensor(self.patch_scale_features.shape[-1]) self.sh_embed_dim_t = torch.tensor(self.sh_embed_dim) self.harmonic_order_t = torch.tensor(self.sh_order_L) @@ -251,8 +252,8 @@ def _compute_geoscalepatch_embeddings(self): patch_geo_embeddings.append(geo_emb) patch_scale_features.append(scale_feat) - patch_geo_embeddings = torch.stack(patch_geo_embeddings) - patch_scale_features = torch.stack(patch_scale_features) + patch_geo_embeddings = torch.stack(patch_geo_embeddings).contiguous().clone() + patch_scale_features = torch.stack(patch_scale_features).contiguous().clone() return patch_geo_embeddings, patch_scale_features diff --git a/climanet/geo_embedding_utils.py b/climanet/geo_embedding_utils.py index 7a1e905..b446421 100644 --- a/climanet/geo_embedding_utils.py +++ b/climanet/geo_embedding_utils.py @@ -49,8 +49,8 @@ def compute_sh_on_grid(lat, lon, L, dtype=torch.float32): H = len(lat) W = len(lon) - theta = torch.deg2rad(90.0 - torch.tensor(lat, dtype=dtype)) # colatitude - phi = torch.deg2rad(torch.tensor(lon, dtype=dtype)) # longitude + theta = torch.deg2rad(90.0 - torch.as_tensor(lat, dtype=dtype)) # colatitude + phi = torch.deg2rad(torch.as_tensor(lon, dtype=dtype)) # longitude D = (L + 1) ** 2 @@ -90,7 +90,7 @@ def compute_area_weights(lat): weights : (H,) """ - lat_rad = torch.deg2rad(torch.tensor(lat, dtype=torch.float32)) + lat_rad = torch.deg2rad(torch.as_tensor(lat, dtype=torch.float32)) weights = torch.cos(lat_rad) diff --git a/climanet/predict.py b/climanet/predict.py index 0ad5c23..ec5dbd4 100644 --- a/climanet/predict.py +++ b/climanet/predict.py @@ -1,10 +1,11 @@ +from pathlib import Path + import numpy as np from torch.utils.data import Dataset -from climanet.st_encoder_decoder import SpatioTemporalModel import xarray as xr import torch from torch.utils.data import DataLoader -from climanet.utils import setup_logging, compute_masked_loss +from climanet.utils import load_model, setup_logging, compute_masked_loss def _save_netcdf(predictions: np.ndarray, dataset: Dataset, save_dir: str): @@ -47,14 +48,6 @@ def _save_netcdf(predictions: np.ndarray, dataset: Dataset, save_dir: str): return ds_pred -def _load_model(model_path: str, device: str): - """Helper function to load a model from a checkpoint.""" - checkpoint = torch.load(model_path, map_location=device, weights_only=False) - model = SpatioTemporalModel(**checkpoint["model_config"]) - model.load_state_dict(checkpoint["model_state_dict"]) - return model.to(device) - - def predict_monthly_var( model: torch.nn.Module | str, dataset: Dataset, @@ -89,8 +82,8 @@ def predict_monthly_var( If return_loss is True, it also returns the average loss over the dataset. """ # Load the model if a path is provided - if isinstance(model, str): - model = _load_model(model, device) + if isinstance(model, str | Path): + model = load_model(model, device) model.to(device) model.eval() diff --git a/climanet/st_encoder_decoder.py b/climanet/st_encoder_decoder.py index 70e9901..db6ad6a 100644 --- a/climanet/st_encoder_decoder.py +++ b/climanet/st_encoder_decoder.py @@ -244,7 +244,7 @@ class TemporalAttentionAggregator(nn.Module): months. """ - def __init__(self, embed_dim=128, dropout=0.0): + def __init__(self, embed_dim=128, dropout=0.0, chunk_size=1024): """Initialize the temporal attention aggregator. Args: @@ -253,10 +253,14 @@ def __init__(self, embed_dim=128, dropout=0.0): of 64 (e.g., 64, 128, 256). This can be tuned. dropout: Dropout rate for regularization in the day scorer and cross-month mixing. Default is 0.0. Increase it if there is overfitting. + chunk_size: Number of chunks to process for memory efficiency. This + is related to PyTorch limitation in PyTorch's efficient attention + kernels. Default is 1024. """ super().__init__() self.time_embed = CyclicTimeEmbedding(embed_dim=embed_dim) + self.chunk_size = chunk_size # cyclical embedding for months self.month_embed = CyclicMonthEmbedding(embed_dim=embed_dim) @@ -288,6 +292,18 @@ def __init__(self, embed_dim=128, dropout=0.0): nn.Linear(4 * embed_dim, embed_dim), ) + def _chunk_attn(self, z): + """Chunked attention to reduce memory usage for long sequences. + + Args: + z: Input tensor of shape (B*HW, M, C) where M is the number of months + """ + out_chunks = [] + for z_chunk in z.split(self.chunk_size, dim=0): + attn_out, _ = self.month_attn(z_chunk, z_chunk, z_chunk, need_weights=False) + out_chunks.append(attn_out) + return torch.cat(out_chunks, dim=0) # (B*HW, M, C) + def forward(self, x, M, time_features, padded_days_mask=None): """ Args: @@ -337,7 +353,7 @@ def forward(self, x, M, time_features, padded_days_mask=None): z = self.month_ln(z) - attn_out, _ = self.month_attn(z, z, z, need_weights=False, is_causal=False) + attn_out = self._chunk_attn(z) z = z + attn_out + self.month_ffn(z) z = z.reshape(B, HW, M, C) diff --git a/climanet/train.py b/climanet/train.py index 4e95fa4..589da65 100644 --- a/climanet/train.py +++ b/climanet/train.py @@ -1,7 +1,9 @@ +from ray import tune from torch.utils.data import Dataset from torch.optim.lr_scheduler import ReduceLROnPlateau from torch.utils.data import DataLoader import torch +from pathlib import Path from climanet.predict import predict_monthly_var from climanet.utils import setup_logging, compute_masked_loss, save_model @@ -38,6 +40,7 @@ def train_monthly_model( verbose: bool = True, dataloader_num_workers: int = 2, verbose_epoch_interval: int = 20, + tune_checkpoint: bool = False, ): """Train the model to predict monthly data from daily data. Args: @@ -82,6 +85,14 @@ def train_monthly_model( counter = 0 best_state_dict = None # Store best model state + if tune.get_checkpoint(): + loaded_checkpoint = tune.get_checkpoint() + with loaded_checkpoint.as_directory() as loaded_checkpoint_dir: + loaded_checkpoint_dir = Path(loaded_checkpoint_dir).resolve() + checkpoint = torch.load(loaded_checkpoint_dir / "checkpoint.pt") + model.load_state_dict(checkpoint["model_state_dict"]) + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + # Add scheduler - reduces LR instead of stopping immediately scheduler = ReduceLROnPlateau( optimizer, @@ -175,6 +186,13 @@ def train_monthly_model( if best_state_dict is not None: model.load_state_dict(best_state_dict) + if tune_checkpoint: + # Save the model and optimizer state for Ray Tune checkpointing + save_model(model, optimizer, run_dir, filename="checkpoint.pt", verbose=False) + tune.report( + {"loss": best_loss}, checkpoint=tune.Checkpoint.from_directory(run_dir) + ) + # Close the writer when done writer.close() diff --git a/climanet/tune.py b/climanet/tune.py new file mode 100644 index 0000000..df769b6 --- /dev/null +++ b/climanet/tune.py @@ -0,0 +1,163 @@ +from pathlib import Path + +import ray +import xarray as xr + +from ray.tune.schedulers import ASHAScheduler +from climanet.dataset import STDataset +from climanet.predict import predict_monthly_var +from climanet.st_encoder_decoder import SpatioTemporalModel +from climanet.train import train_monthly_model +from climanet.utils import set_seed + + +def _train(tune_config): + """Helper function to train the model with Ray Tune.""" + + device = tune_config["device"] + dataloader_num_workers = tune_config["dataloader_num_workers"] + train_dataset = ray.get(tune_config["train_dataset"]) + validation_dataset = ray.get(tune_config["validation_dataset"]) + patch_size = tune_config["patch_size"] + overlap = tune_config["overlap"] + embed_dim = tune_config["embed_dim"] + dropout = tune_config["dropout"] + hidden = tune_config["hidden"] + spatial_depth = tune_config["spatial_depth"] + spatial_heads = tune_config["spatial_heads"] + run_dir = tune_config["run_dir"] + num_epoch = tune_config["num_epoch"] + + set_seed() + + model = SpatioTemporalModel( + patch_size=(1, patch_size, patch_size), + overlap=overlap, + embed_dim=embed_dim, + dropout=dropout, + hidden=hidden, + spatial_depth=spatial_depth, + spatial_heads=spatial_heads, + ) + + batch_size = tune_config["batch_size"] + accumulation_steps = tune_config["accumulation_steps"] + optimizer_lr = tune_config["optimizer_lr"] + + _ = train_monthly_model( + model, + train_dataset, + validation_dataset=validation_dataset, + batch_size=batch_size, + num_epoch=num_epoch, + accumulation_steps=accumulation_steps, + optimizer_lr=optimizer_lr, + device=device, + run_dir=run_dir, + dataloader_num_workers=dataloader_num_workers, + store_model=False, + verbose=False, + tune_checkpoint=True, + ) + + +def tune_data_preparation(data_config: dict, is_hourly=True) -> STDataset: + """Prepare the data for training and validation.""" + input_data = xr.open_mfdataset(data_config["input_filenames"], chunks=data_config.get("input_chunks")) + monthly_data = xr.open_mfdataset(data_config["monthly_filenames"], chunks=data_config.get("monthly_chunks")) + lsm_mask = xr.open_dataset(data_config["landmask_filename"], chunks=data_config.get("landmask_chunks")) + + # calculate residuals as target + input_data_averaged = input_data.resample(time="MS").mean(skipna=True) + input_data_averaged["time"] = monthly_data["time"] + + # Residuals + monthly_data_res = monthly_data - input_data_averaged + + var_name = data_config["var_name"] + + dataset = STDataset( + input_da=input_data[var_name], + monthly_da=monthly_data_res[var_name], + land_mask=lsm_mask["lsm"], + patch_size=data_config["patch_size"], + stride=data_config["stride"], + sh_embed_dim=96, + sh_order_L=10, + is_hourly=is_hourly, + ) + return dataset + + +def run_tune(tune_config: dict): + """Run Ray Tune to find the best hyperparameters for the model. + + Args: + tune_config: dictionary containing the hyperparameters to tune and their + ranges and other config parameters. + """ + scheduler = ASHAScheduler( + time_attr="training_iteration", + max_t=tune_config["max_num_epochs"], + grace_period=1, + reduction_factor=2, + ) + + tuner = ray.tune.Tuner( + ray.tune.with_resources( + ray.tune.with_parameters(_train), + resources={ + "cpu": tune_config["cpu_per_trial"], + "gpu": tune_config["gpu_per_trial"], + }, + ), + tune_config=ray.tune.TuneConfig( + metric="loss", + mode="min", + scheduler=scheduler, + num_samples=tune_config["num_trials"], + max_concurrent_trials=tune_config["max_concurrent_trials"], + ), + param_space=tune_config, + run_config=ray.tune.RunConfig(storage_path=tune_config["run_dir"]), + ) + results = tuner.fit() + return results + + +def check_best_model(experiment_path: str | Path, test_dataset: STDataset, run_dir: str | Path): + """Test the best model from a Ray Tune experiment. + + Args: + experiment_path: path to the Ray Tune experiment directory + test_dataset: Dataset object containing the test data + run_dir: directory to save logs and model + Returns: + test_loss: the loss on the test dataset + + """ + if not ray.is_initialized(): + ray.init() + analysis = ray.tune.ExperimentAnalysis(experiment_path) + best_result = analysis.get_best_trial("loss", "min") + + batch_size = best_result.config["batch_size"] + device = best_result.config["device"] + dataloader_num_workers = best_result.config["dataloader_num_workers"] + best_checkpoint = best_result.checkpoint + + model_path = Path(best_checkpoint.path) / "checkpoint.pt" + + _, test_loss = predict_monthly_var( + model_path, + test_dataset, + batch_size=batch_size, + device=device, + return_numpy=False, + save_predictions=False, + return_loss=True, + verbose=False, + run_dir=run_dir, + dataloader_num_workers=dataloader_num_workers, + ) + return test_loss diff --git a/climanet/utils.py b/climanet/utils.py index a56ae4a..469c11d 100644 --- a/climanet/utils.py +++ b/climanet/utils.py @@ -12,6 +12,8 @@ from matplotlib.colors import TwoSlopeNorm import matplotlib.ticker as mticker +from climanet.st_encoder_decoder import SpatioTemporalModel + def regrid_to_boundary_centered_grid(da: xr.DataArray, roll=False) -> xr.DataArray: """ @@ -254,18 +256,37 @@ def compute_masked_loss( return (num / denom).mean() -def save_model(model: torch.nn.Module, run_dir: str, verbose: bool) -> None: +def save_model( + model: torch.nn.Module, + optimizer: torch.optim.Optimizer, + run_dir: str, + filename="best_model.pth", + verbose: bool = True, +) -> None: """Save model state and config to disk.""" - model_path = Path(run_dir) / "best_model.pth" + Path(run_dir).mkdir(parents=True, exist_ok=True) + model_path = Path(run_dir) / filename model = model.module if hasattr(model, "module") else model torch.save( - {"model_state_dict": model.state_dict(), "model_config": model.config}, + { + "model_state_dict": model.state_dict(), + "optimizer_state_dict": optimizer.state_dict(), + "model_config": model.config, + }, model_path, ) if verbose: print(f"Model saved to {model_path}") +def load_model(model_path: str, device: str): + """Helper function to load a model from a checkpoint.""" + checkpoint = torch.load(model_path, map_location=device, weights_only=False) + model = SpatioTemporalModel(**checkpoint["model_config"]) + model.load_state_dict(checkpoint["model_state_dict"]) + return model.to(device) + + def add_month_hour_dims( hourly_ts: xr.DataArray, # (time, H, W) hourly monthly_ts: xr.DataArray, # (time, H, W) monthly @@ -468,6 +489,40 @@ def plot_histograms( plt.show() +def data_split( + data_folder, + filename_pattern="*_hr_ERA5dc_masked_tos.nc", + train_range=(2018, 2020), + validation_range=(2021, 2021), + test_range=(2022, 2022), +): + """ + Split the data into training, validation, and test sets based on the provided year ranges. + """ + data_folder = Path(data_folder) + + splits = { + "train": [], + "validation": [], + "test": [], + } + + for file in data_folder.rglob(filename_pattern): + year = int(file.stem[:4]) + + if train_range[0] <= year <= train_range[1]: + splits["train"].append(file) + elif validation_range[0] <= year <= validation_range[1]: + splits["validation"].append(file) + elif test_range[0] <= year <= test_range[1]: + splits["test"].append(file) + + for lst in splits.values(): + lst.sort() + + return splits + + def plot_nobs_vs_err( nobs: xr.DataArray, err_baseline: xr.DataArray, err_predictions: xr.DataArray ): @@ -487,7 +542,7 @@ def plot_nobs_vs_err( for i, ax in enumerate(axes): ax.set_title(f"Month = {err_baseline.time.dt.strftime('%Y-%m-%d').values[i]}") - + # Get unique number of observations for this month, ignoring NaNs and zeros n_obs_unique = np.unique(nobs.isel(time=i).values) n_obs_unique = n_obs_unique[(~np.isnan(n_obs_unique)) & (n_obs_unique > 0)] diff --git a/notebooks/example_tuning.ipynb b/notebooks/example_tuning.ipynb new file mode 100644 index 0000000..5a388b2 --- /dev/null +++ b/notebooks/example_tuning.ipynb @@ -0,0 +1,400 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "57546057-2042-42eb-a793-23e77d22965e", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2026-07-16 17:16:54,174\tINFO util.py:155 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n", + "2026-07-16 17:16:56,404\tINFO util.py:155 -- Missing packages: ['ipywidgets']. Run `pip install -U ipywidgets`, then restart the notebook server for rich notebook output.\n" + ] + } + ], + "source": [ + "from pathlib import Path\n", + "\n", + "import ray\n", + "import torch\n", + "\n", + "from climanet.tune import run_tune, check_best_model\n", + "from climanet.dataset import STDataset\n", + "from torch.utils.data import random_split\n", + "\n", + "import xarray as xr" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "d9307dee-167b-4d7f-93d4-f0525a449baa", + "metadata": {}, + "outputs": [], + "source": [ + "data_folder = Path(\"./eso4clima/dc_data\")\n", + "run_dir = Path(\"./runs_daily\").resolve()\n", + "\n", + "var_name = \"tos\"\n", + "input_filenames = [data_folder / f\"202101_day_ERA5dc_masked_{var_name}.nc\", data_folder / f\"202102_day_ERA5dc_masked_{var_name}.nc\"]\n", + "monthly_filenames = [data_folder / f\"202101_mon_ERA5dc_full_{var_name}.nc\", data_folder / f\"202102_mon_ERA5dc_full_{var_name}.nc\"]" + ] + }, + { + "cell_type": "markdown", + "id": "64f35b02-6c72-4393-8aa6-5b23fa46d6ee", + "metadata": {}, + "source": [ + "### Load the data" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "acb3cf47-2b92-439a-9116-e1d3e366e4ff", + "metadata": {}, + "outputs": [], + "source": [ + "data_config = {\n", + " \"input_filenames\": input_filenames,\n", + " \"monthly_filenames\": monthly_filenames,\n", + " \"landmask_filename\": data_folder / \"era5_lsm_bool.nc\",\n", + " \"var_name\": var_name,\n", + " \"patch_size\": (1, 40, 40), # based on the patch_size in model\n", + " \"stride\": (20, 20), # data agumentation by overlapping patches\n", + "}" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "8236eb50-9a0b-4cc2-995c-3d84e2db9f3a", + "metadata": {}, + "outputs": [], + "source": [ + "input_data = xr.open_mfdataset(data_config[\"input_filenames\"])\n", + "monthly_data = xr.open_mfdataset(data_config[\"monthly_filenames\"])\n", + "lsm_mask = xr.open_dataset(data_config[\"landmask_filename\"])" + ] + }, + { + "cell_type": "markdown", + "id": "3fec14cc-97a8-4e95-b6c7-4ffe7e476693", + "metadata": {}, + "source": [ + "### prepare data for tuning" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "0eaf7500-b1e3-48a1-bd20-25d7f0b492a1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Patch grid (m x i x j): 2 x 7 x 19 = 266 patches\n", + "Overlap: 20 pixels (height), 20 pixels (width)\n" + ] + } + ], + "source": [ + "# coordinates of subset\n", + "lon_subset = slice(-50, 50) # one lon -179.9 is nan, check data\n", + "lat_subset = slice(-30, 10)\n", + "\n", + "input_data = input_data.sel(lon=lon_subset, lat=lat_subset)\n", + "monthly_data = monthly_data.sel(lon=lon_subset, lat=lat_subset)\n", + "lsm_mask = lsm_mask.sel(lon=lon_subset, lat=lat_subset) \n", + "\n", + "# calculate residuals as target\n", + "input_data_averaged = input_data.resample(time=\"MS\").mean(skipna=True)\n", + "input_data_averaged[\"time\"] = monthly_data[\"time\"]\n", + "\n", + "# Residuals\n", + "monthly_data_res = monthly_data - input_data_averaged\n", + "\n", + "var_name = data_config[\"var_name\"]\n", + "\n", + "dataset = STDataset(\n", + " input_da=input_data[var_name],\n", + " monthly_da=monthly_data_res[var_name],\n", + " land_mask=lsm_mask[\"lsm\"],\n", + " patch_size=data_config[\"patch_size\"],\n", + " stride=data_config[\"stride\"],\n", + " sh_embed_dim=96,\n", + " sh_order_L=10,\n", + " is_hourly=False,\n", + "\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "babb4125-33dd-4df4-a12d-22041634fc98", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "159 79 28\n" + ] + } + ], + "source": [ + "# create train test data\n", + "generator = torch.Generator().manual_seed(42)\n", + "train_size = int(0.6 * len(dataset))\n", + "validation_size = int(0.3 * len(dataset))\n", + "test_size = len(dataset) - train_size - validation_size\n", + "train_dataset, validation_dataset, test_dataset = random_split(dataset, [train_size, validation_size, test_size], generator=generator)\n", + "print(len(train_dataset), len(validation_dataset), len(test_dataset))" + ] + }, + { + "cell_type": "markdown", + "id": "c0b02899-0208-4949-a152-3f5ce27451fc", + "metadata": {}, + "source": [ + "### config for hyper parameter tuning" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "f1ebcf7b-61d0-4f7c-a895-91ffce764f58", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2026-07-16 17:17:08,413\tINFO worker.py:2024 -- Started a local Ray instance.\n" + ] + } + ], + "source": [ + "tune_config = {\n", + " \"max_num_epochs\": 1,\n", + " \"num_trials\": 1,\n", + " \"cpu_per_trial\": 2,\n", + " \"gpu_per_trial\": 0,\n", + " \"run_dir\": run_dir,\n", + " \"device\": \"cpu\",\n", + " \"dataloader_num_workers\": 2,\n", + " \"train_dataset\": ray.put(train_dataset),\n", + " \"validation_dataset\": ray.put(validation_dataset),\n", + " \"num_epoch\": 1,\n", + " # parameters to tune\n", + " \"patch_size\": ray.tune.grid_search([2, 4]),\n", + " \"overlap\": 2,\n", + " \"embed_dim\": 32,\n", + " \"dropout\": 0.0,\n", + " \"hidden\": 32,\n", + " \"spatial_depth\": 3,\n", + " \"spatial_heads\": 4,\n", + " \"optimizer_lr\": 1e-1,\n", + " \"batch_size\": 10,\n", + " \"accumulation_steps\": 5,\n", + " \"max_concurrent_trials\": 2\n", + "}" + ] + }, + { + "cell_type": "markdown", + "id": "c619d9a4-d388-4a47-9cac-04e773dc4657", + "metadata": {}, + "source": [ + "### Run ray tune" + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "be79c19a-5673-4672-b9d9-b1d93da4b0b7", + "metadata": {}, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "
\n", + "
\n", + "

Tune Status

\n", + " \n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
Current time:2026-07-16 16:50:12
Running for: 00:01:13.04
Memory: 9.0/15.3 GiB
\n", + "
\n", + "
\n", + "
\n", + "

System Info

\n", + " Using AsyncHyperBand: num_stopped=2
Bracket: Iter 1.000: -0.37507010996341705
Logical resource usage: 2.0/8 CPUs, 0/0 GPUs\n", + "
\n", + " \n", + "
\n", + "
\n", + "
\n", + "

Trial Status

\n", + " \n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "\n", + "
Trial name status loc patch_size iter total time (s) loss
_train_7a5a1_00000TERMINATED192.168.2.13:392223 2 1 67.64370.384662
_train_7a5a1_00001TERMINATED192.168.2.13:392222 4 1 65.14120.365478
\n", + "
\n", + "
\n", + "\n" + ], + "text/plain": [ + "" + ] + }, + "metadata": {}, + "output_type": "display_data" + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2026-07-16 16:50:12,244\tINFO tune.py:1007 -- Wrote the latest version of all result files and experiment state to '/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-16_16-48-59' in 0.0057s.\n", + "2026-07-16 16:50:12,250\tINFO tune.py:1039 -- Total run time: 73.09 seconds (73.04 seconds for the tuning loop).\n", + "\u001b[36m(_train pid=392223)\u001b[0m Checkpoint successfully created at: Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-16_16-48-59/_train_7a5a1_00000_0_patch_size=2_2026-07-16_16-48-59/checkpoint_000000)\n" + ] + } + ], + "source": [ + "results = run_tune(tune_config)\n", + "\n", + "ray.shutdown()" + ] + }, + { + "cell_type": "markdown", + "id": "db963822-fd41-49f8-8e7f-0cf0d3089a9d", + "metadata": {}, + "source": [ + "### Inspect the output" + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "2913a2b4-adf4-42f1-9f2a-507fd828ffe8", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "Result(\n", + " metrics={'loss': 0.36547836661338806},\n", + " path='/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-16_16-48-59/_train_7a5a1_00001_1_patch_size=4_2026-07-16_16-48-59',\n", + " filesystem='local',\n", + " checkpoint=Checkpoint(filesystem=local, path=/home/sarah/GitHub/ClimaNet/notebooks/runs_daily/_train_2026-07-16_16-48-59/_train_7a5a1_00001_1_patch_size=4_2026-07-16_16-48-59/checkpoint_000000)\n", + ")" + ] + }, + "execution_count": 14, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "results.get_best_result(\"loss\", \"min\")" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "1026753d-9b86-4f34-8b18-d7753cf0dabc", + "metadata": {}, + "outputs": [ + { + "data": { + "text/plain": [ + "0.4014973243077596" + ] + }, + "execution_count": 8, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "experiment_path = results.experiment_path # or add the path above manually as Path(\"./runs_daily/_train_2026-07-16_16-48-59\").resolve()\n", + "check_best_model(experiment_path, test_dataset, run_dir)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fef22a4c-5df2-47e7-806e-f1edd0a3fa37", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.11.14" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index e83d4b3..10c9798 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "tensorboard", "torch>=2.10.0", "xarray>=2025.12.0", + "ray[tune]", ] [dependency-groups] diff --git a/scripts/tuning.py b/scripts/tuning.py new file mode 100644 index 0000000..d0011b6 --- /dev/null +++ b/scripts/tuning.py @@ -0,0 +1,102 @@ +import argparse +from pathlib import Path + +from ray import tune +import ray + +from climanet.tune import run_tune, tune_data_preparation +from climanet.utils import data_split + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument( + "--storage-path", + type=str, + default=Path("./tune_results").resolve(), + ) + parser.add_argument( + "--num-nodes", + type=int, + default=1, + ) + parser.add_argument( + "--data-dir", + type=str, + default=Path("./data").resolve(), + ) + parser.add_argument( + "--lsm-dir", + type=str, + default=Path("./data").resolve(), + ) + parser.add_argument( + "--ray-address", + type=str, + default="auto", + ) + args = parser.parse_args() + + data_folder = Path(args.data_dir).resolve() + lsm_folder = Path(args.lsm_dir).resolve() + hourly_files = data_split( + data_folder, + filename_pattern="*_hr_ERA5dc_masked_tos.nc", + train_range=(2018, 2020), + validation_range=(2021, 2021), + test_range=(2022, 2022), + ) + monthly_files = data_split( + data_folder, + filename_pattern="*_mon_ERA5dc_masked_tos.nc", + train_range=(2018, 2020), + validation_range=(2021, 2021), + test_range=(2022, 2022), + ) + data_config_train = { + "input_filenames": hourly_files["train"], + "monthly_filenames": monthly_files["train"], + "landmask_filename": lsm_folder / "era5_lsm_bool.nc", + "var_name": "tos", + "patch_size": (1, 40, 40), # based on the patch_size in model + "stride": (20, 20), # data agumentation by overlapping patches + } + data_config_validation = { + "input_filenames": hourly_files["validation"], + "monthly_filenames": monthly_files["validation"], + "landmask_filename": lsm_folder / "era5_lsm_bool.nc", + "var_name": "tos", + "patch_size": (1, 40, 40), # based on the patch_size in model + "stride": (20, 20), # data agumentation by overlapping patches + } + + tune_config = { + "max_num_epochs": 100, + "num_trials": 10, + "cpu_per_trial": 4, + "gpu_per_trial": 1, + "run_dir": args.storage_path, + "device": "cuda", + "dataloader_num_workers": 4, + "train_dataset": tune_data_preparation(data_config_train), + "validation_dataset": tune_data_preparation(data_config_validation), + "num_epoch": 100, + # parameters to tune + "patch_size": tune.grid_search([2, 4, 8]), + "overlap": tune.grid_search([0, 1, 2]), + "embed_dim": tune.grid_search([32, 64, 128]), + "dropout": tune.grid_search([0.0, 0.1, 0.2]), + "hidden": tune.grid_search([32, 64, 128]), + "spatial_depth": tune.grid_search([1, 2, 3]), + "spatial_heads": tune.grid_search([2, 4, 8]), + "optimizer_lr": tune.loguniform(1e-3, 1e-1), + "batch_size": tune.grid_search([200, 400, 800]), # based on GPU memory + "accumulation_steps": tune.grid_search([1, 2, 4]), # based on batch_size + "max_concurrent_trials": args.num_nodes * 4, # GPUs per node (4) + } + + # Start Ray Tune for distributed training on several nodes + ray.init(address=args.ray_address, ignore_reinit_error=True) + + results = run_tune(tune_config) + + ray.shutdown() diff --git a/scripts/tuning.slurm b/scripts/tuning.slurm new file mode 100644 index 0000000..058cea0 --- /dev/null +++ b/scripts/tuning.slurm @@ -0,0 +1,91 @@ +#!/bin/bash +#SBATCH --job-name=tuning +#SBATCH --partition=gpu +#SBATCH --nodes=10 +#SBATCH --ntasks-per-node=1 +#SBATCH --gpus-per-task=4 +#SBATCH --time=04:00:00 +#SBATCH --account=bd0854 +#SBATCH --output=tuning_%j.out + +set -euo pipefail +ulimit -s 204800 + +# Activate uv env +UV_ENV="$HOME/.venv" +source "$UV_ENV/bin/activate" + +# Set the scratch directory because they are avialable to all nodes +SCRATCH="/scratch/b/$USER/tune_results" + +# data directory (adjust this path to your data location) +DATA_DIR="$HOME/scratch/b/$USER/data" +LSM_DIR="$HOME/scratch/b/$USER/data" + +# Get node list +nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") + +# First node is the Ray head +head_node="${nodes[0]}" +head_node_ip=$(srun --nodes=1 --ntasks=1 -w "$head_node" hostname -I | awk '{print $1}') + +echo "Ray head: $head_node_ip" + +# Start Ray head on the first node +# 6379 default port for Ray’s Redis/GCS service +srun --nodes=1 --ntasks=1 -w "$head_node" \ + bash -lc " + source '$UV_ENV/bin/activate' && + export OMP_NUM_THREADS=1 && + ray start --head \ + --node-ip-address=$head_node_ip \ + --port=6379 \ + --num-cpus=$SLURM_CPUS_PER_TASK \ + --num-gpus=4 \ + --block + " & + +sleep 15 + +# Start Ray workers on the remaining nodes +for node in "${nodes[@]:1}"; do + srun --nodes=1 --ntasks=1 -w "$node" \ + bash -lc " + source '$UV_ENV/bin/activate' && + export OMP_NUM_THREADS=1 && + ray start --address=$head_node_ip:6379 \ + --num-cpus=$SLURM_CPUS_PER_TASK \ + --num-gpus=4 \ + --block + " & +done + +sleep 15 + +echo "=== Ray sanity check ===" +srun --nodes=1 --ntasks=1 -w "$head_node" \ + bash -lc "source '$UV_ENV/bin/activate' && python - <<'PY' +import os +import ray + +address = os.environ.get('RAY_ADDRESS', '${head_node_ip}:6379') +ray.init(address=address, ignore_reinit_error=True) + +resources = ray.available_resources() +print('Address:', address) +print('CPU:', resources.get('CPU')) +print('GPU:', resources.get('GPU')) +print('Resources:', resources) + +ray.shutdown() +PY" + +sleep 15 + +# Run the tuning script on the head node +python -u ~/scripts/tuning.py \ + --storage-path "$SCRATCH" \ + --num-nodes "$SLURM_NNODES" \ + --ray-address "$head_node_ip:6379" \ + --data-dir "$DATA_DIR" \ + --lsm-dir "$LSM_DIR"