Fix generalization

climanet/dataset.py

@@ -44,6 +44,7 @@ def __init__(
         monthly_da: xr.DataArray,
         padded_days_mask: xr.DataArray,
         time_features: xr.DataArray,
+        res_std: xr.DataArray,
         land_mask: xr.DataArray = None,
         model_patch_size: tuple[int, int, int] = (1, 4, 4),  # (Month, lat, lon)
         spatial_dims: tuple[str, str] = ("lat", "lon"),
@@ -151,6 +152,7 @@ def __init__(
         self.geo_pos_t = calculate_sh_geo_pos_embeddings(
             self.lat_coords, self.lon_coords, self.sh_order_L, self.sh_embed_dim
         )
+        self.res_std_t = torch.from_numpy(res_std.to_numpy()).contiguous()
 
     def _compute_crop_indices(self, M: int, H: int, W: int) -> list:
         """Generate crop start indices with coverage warning (overlap support)."""
@@ -322,6 +324,7 @@ def __getitem__(self, idx):
         # compute geo_pos_embedding and scale_feature for each model patch within the dataset crop
         geo_pos_embedding_t = self._compute_patch_geo_pos_embedding(i, j)  # (Hp*Wp, sh_embed_dim)
         scale_feature_t = self._compute_patch_scale_features(i, j)  # (Hp*Wp, scale_dim)
+        res_std_t = self.res_std_t[m : m + pm, i : i + ph, j : j + pw]
 
         # Convert to dictionary
         return {
@@ -338,6 +341,7 @@ def __getitem__(self, idx):
             "coords": torch.tensor([m, i, j]),
             "lat": lat_crop,  # (pH,)
             "lon": lon_crop,  # (pW,)
+            "res_std": res_std_t,  # (pm, pH, pW)
         }
 
     def __getitems__(self, indices):

climanet/predict.py

@@ -1,17 +1,17 @@
 from dataclasses import dataclass
 from pathlib import Path
 
 import numpy as np
 import torch
 import xarray as xr
 from torch.utils.data import DataLoader, Dataset
 
 from climanet.dataset import DataLoaderConfig
 from climanet.utils import (
-    compute_masked_loss,
+    run_one_batch,
     load_model,
     setup_logging,
 )
 
 
 @dataclass
@@ -74,28 +74,6 @@
         ds_pred.sel(time=[t]).to_netcdf(file_name)
 
 
-def _move_batch_to_device(batch: dict, device: str):
-    use_cuda = device == "cuda"
-    return {k: v.to(device, non_blocking=use_cuda) for k, v in batch.items()}
-
-
-def _run_one_batch(model: torch.nn.Module, batch: dict, device: str):
-    batch = _move_batch_to_device(batch, device)
-    pred = model(
-        batch["input_data"],
-        batch["input_data_mask"],
-        batch["input_data_timef"],
-        batch["land_mask"],
-        batch["geo_pos_embedding"],
-        batch["scale_feature"],
-        batch["padded_days_mask"],
-    )  # (B, M, H, W)
-
-    # Compute masked loss
-    loss = compute_masked_loss(pred, batch["monthly_data"], batch["land_mask"])
-    return loss, pred
-
-
 def predict_monthly_var(
     model: torch.nn.Module | str,
     dataset: Dataset,
@@ -132,7 +110,7 @@
     dataloader = DataLoader(
         dataset,
         batch_size=dataloader_config.batch_size,
-        shuffle=dataloader_config.shuffle,
+        shuffle=False,  # no shuffling during prediction/validation
         pin_memory=use_cuda,
         num_workers=dataloader_config.num_workers,  # for data loading
         persistent_workers=dataloader_config.persistent_workers,  # keep workers alive between epochs
@@ -158,7 +136,7 @@
         idx = 0
         average_loss = 0.0
         for i, batch in enumerate(dataloader):
-            loss, predictions = _run_one_batch(model, batch, device)
+            loss, predictions = run_one_batch(model, batch, device)
             average_loss += loss.detach()
 
             all_predictions[idx : idx + predictions.size(0)] = predictions.detach()

climanet/st_encoder_decoder.py

@@ -3,9 +3,10 @@
 The main model class is SpatioTemporalModel.
 """
 
 import torch
 from torch import nn
 from torch.utils.checkpoint import checkpoint
+import math
 
 
 class VideoEncoder(nn.Module):
@@ -41,6 +42,7 @@
             2 * in_chans, embed_dim, kernel_size=patch_size, stride=patch_size
         )
 
+
     def forward(self, x, mask):
         """Forward pass with masking support via an additional validity channel.
         Args:
@@ -66,7 +68,6 @@
         x = x.contiguous(memory_format=torch.channels_last_3d)
         x = self.proj(x)
         x = x.flatten(2).transpose(1, 2)
-
         return x
 
 
@@ -381,6 +382,7 @@
         patch_w=4,
         hidden=128,
         dropout=0.0,
+        residual_scale=0.5,
     ):
         """
         Args:
@@ -392,46 +394,69 @@
             hidden: Hidden dimension in the decoder for mixing channel features.
                 The default is 128, which can be tuned.
             dropout: Dropout rate for regularization in the refinement block. Default is 0.0.
+            residual_scale: Scaling factor for the residual connection in the refinement block. Default is 0.5.
         """
         super().__init__()
         self.patch_h = patch_h
         self.patch_w = patch_w
+        self.residual_scale = residual_scale
+
+        # Number of progressive x2 upsampling stages.
+        n_upsample = int(math.log2(patch_h))
+
 
         # Mix channel features on the patch grid (Hp, Wp)
         # Input shape: (B, embed_dim, Hp, Wp) → Output shape: (B, hidden, Hp, Wp)
         # here kernel_size=1 means we are mixing features at each patch location
         # without spatial interaction
         in_channels, out_channels = embed_dim, hidden
-        self.proj = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+        self.proj = nn.Sequential(
+            nn.Conv2d(embed_dim, hidden, kernel_size=1),
+            nn.GroupNorm(8, hidden),
+            nn.GELU(),
+        )
 
-        # Upsample to full resolution
-        in_channels, out_channels = hidden, hidden // 2
-        self.deconv = nn.ConvTranspose2d(
-            in_channels,
-            out_channels,
-            kernel_size=(patch_h, patch_w),
-            stride=(patch_h, patch_w),
-            padding=0,
+        # Spatial processing before upsampling
+        self.spatial_mix = nn.Sequential(
+            nn.Conv2d(hidden, hidden, kernel_size=3, padding=1),
+            nn.GroupNorm(8, hidden),
+            nn.GELU(),
         )
 
+        # Upsample to full resolution
+        channels = hidden
+        self.deconv = nn.ModuleList()
+
+        for i in range(n_upsample):
+            # Gradually reduce channels as spatial resolution increases
+            next_channels = max(hidden // (2 ** (i + 1)), 16)
+            block = nn.Sequential(
+                nn.Upsample(scale_factor=2, mode="bilinear", align_corners=False),
+                nn.Conv2d(channels, next_channels, kernel_size=3, padding=1),
+                nn.GroupNorm(min(8, next_channels), next_channels),
+                nn.GELU(),
+                nn.Dropout2d(dropout),
+            )
+            self.deconv.append(block)
+            channels = next_channels
+
         # Final conv head to get single channel output kernel_size=3 is the most
         # common choice for spatial convolutions; it's the smallest kernel that
         # captures spatial context in all directions
-        in_channels, out_channels = hidden // 2, hidden // 2
-
         # Refinement block: a small conv layers to smooth patch boundaries
         self.refine = nn.Sequential(
-            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups=8, num_channels=out_channels),
-            nn.GELU(),
-            nn.Dropout2d(dropout),
-            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
-            nn.GroupNorm(num_groups=8, num_channels=out_channels),
+            nn.Conv2d(channels, channels, kernel_size=3, padding=1),
+            nn.GroupNorm(min(8, channels), channels),
             nn.GELU(),
         )
 
         # Final conv head to map to single-channel output
-        self.head = nn.Conv2d(out_channels, 1, kernel_size=1)
+        self.head = nn.Conv2d(channels, 1, kernel_size=1)
+
+        # Initialize the final conv head weights and biases to zero for stable
+        # training
+        nn.init.zeros_(self.head.weight)
+        nn.init.zeros_(self.head.bias)
 
     def forward(self, latent, M, out_H, out_W, land_mask=None):
         """Reconstruct 2D maps from latent patch tokens.
@@ -459,8 +484,13 @@
         # Apply 1x1 convolution to mix features
         out = self.proj(out)  # (B*M, hidden, Hp, Wp)
 
-        # Use transposed convolution to upsample
-        out = self.deconv(out)  # (B*M, hidden//2, H, W)
+        # Spatial interaction before upsampling
+        residual = self.spatial_mix(out)
+        out = out + residual
+
+        # Progressive x2 upsampling
+        for block in self.deconv:
+            out = block(out)
 
         # Refinement CNN to smooth boundaries
         refine_input = out
@@ -469,6 +499,10 @@
 
         # Apply final conv head to get single channel output
         out = self.head(out)  # (B*M, 1, H, W)
+
+        # Apply residual scaling and non-linearity
+        out = self.residual_scale * torch.tanh(out)
+
         out = out.view(B, M, out_H, out_W)  # (B, M, H, W)
 
         # Mask out land areas if land_mask is provided
@@ -589,6 +623,7 @@
         sh_dim=96,
         scale_dim=10,
         use_checkpoint=True,
+        decoder_residual_scale=0.5,
     ):
         """Initialize the Spatio-Temporal Model.
 
@@ -633,6 +668,7 @@
             patch_w=patch_size[2],
             hidden=hidden,
             dropout=dropout,
+            residual_scale=decoder_residual_scale,
         )
         self.patch_size = patch_size
 
@@ -742,9 +778,10 @@
         # decoder input shape is (B, M*Hp*Wp, C), C: embedding dimension
         # decoder output shape is (B, M, H, W)
         if self.use_checkpoint:
-            monthly_pred = checkpoint(
+            monthly_res_pred = checkpoint(
                 self.decoder, x, M, H, W, land_mask_patch, use_reentrant=False
             )  # (B, M, H, W)
         else:
-            monthly_pred = self.decoder(x, M, H, W, land_mask_patch)  # (B, M, H, W)
-        return monthly_pred
+            monthly_res_pred = self.decoder(x, M, H, W, land_mask_patch)  # (B, M, H, W)
+
+        return monthly_res_pred

climanet/train.py

@@ -1,3 +1,4 @@
+import copy
 import tempfile
 from dataclasses import dataclass
 from pathlib import Path
@@ -10,7 +11,7 @@
 from climanet.dataset import DataLoaderConfig
 from climanet.predict import PredictionConfig, predict_monthly_var
 from climanet.utils import (
-    compute_masked_loss,
+    run_one_batch,
     save_model,
     setup_logging,
 )
@@ -36,27 +37,6 @@ class TrainConfig:
     store_logs: bool = True
 
 
-def _move_batch_to_device(batch: dict, device: str):
-    use_cuda = device == "cuda"
-    return {k: v.to(device, non_blocking=use_cuda) for k, v in batch.items()}
-
-
-def _run_one_batch(model: torch.nn.Module, batch: dict, device):
-    batch = _move_batch_to_device(batch, device)
-    pred = model(
-        batch["input_data"],
-        batch["input_data_mask"],
-        batch["input_data_timef"],
-        batch["land_mask"],
-        batch["geo_pos_embedding"],
-        batch["scale_feature"],
-        batch["padded_days_mask"],
-    )  # (B, M, H, W)
-
-    # Compute masked loss
-    return compute_masked_loss(pred, batch["monthly_data"], batch["land_mask"])
-
-
 def _load_checkpoint(model, optimizer, loaded_checkpoint):
     with loaded_checkpoint.as_directory() as loaded_checkpoint_dir:
         loaded_checkpoint_dir = Path(loaded_checkpoint_dir).resolve()
@@ -136,7 +116,7 @@ def train_monthly_model(
         optimizer.zero_grad()
 
         for i, batch in enumerate(dataloader):
-            loss = _run_one_batch(model, batch, device)
+            loss, _ = run_one_batch(model, batch, device, return_predictions=False)
 
             # Scale loss for gradient accumulation
             scaled_loss = loss * (1.0 / training_config.accumulation_steps)
@@ -196,6 +176,7 @@ def train_monthly_model(
             ):
                 gap = avg_val_loss - avg_train_loss
                 print(f"Epoch {epoch}: gap between train and val loss: {gap:.6f}")
+                print(f"Epoch {epoch}: train loss = {avg_train_loss:.6f}, val loss = {avg_val_loss:.6f}")
 
         # Step scheduler
         scheduler.step(avg_epoch_loss)
@@ -204,7 +185,7 @@ def train_monthly_model(
         # Consider improvement only if loss decreases more than a small threshold
         if avg_epoch_loss < best_loss - 1e-4:
             best_loss = avg_epoch_loss
-            best_state_dict = {k: v.detach() for k, v in model.state_dict().items()}
+            best_state_dict = copy.deepcopy(model.state_dict())
             counter = 0
         else:
             counter += 1

climanet/tune.py

@@ -67,11 +67,12 @@ def _train(tune_config, static_args):
 
     # Prepare training configuration
     training_config = TrainConfig(
-        calculate_residuals=True,
+        calculate_residuals=False,
         num_epoch=static_args["num_epoch"],
         patience=10,
         accumulation_steps=tune_config["batch_config"]["accumulation_steps"],
         optimizer_lr=tune_config["optimizer_lr"],
+        optimizer_weight_decay=tune_config["optimizer_weight_decay"],
         device=static_args["device"],
         verbose=False,
         verbose_epoch_interval=20,
@@ -84,11 +85,13 @@ def _train(tune_config, static_args):
     embed_dim = tune_config["embed_dim"]
     dropout = tune_config["dropout"]
     hidden = tune_config["hidden"]
+    decoder_residual_scale = tune_config["decoder_residual_scale"]
     model = SpatioTemporalModel(
         patch_size=(1, patch_size, patch_size),
         embed_dim=embed_dim,
         dropout=dropout,
         hidden=hidden,
+        decoder_residual_scale=decoder_residual_scale,
     )
 
     # Train the model