Hyperparameter tuning

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
 

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)
 

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()

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)

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()
 

climanet/tune.py

@@ -0,0 +1,161 @@
+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"])
+    monthly_data = xr.open_mfdataset(data_config["monthly_filenames"])
+    lsm_mask = xr.open_dataset(data_config["landmask_filename"])
+
+    # 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, test_dataset: STDataset, run_dir):
+    """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 = f"{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