Fine tuning Aurora model issue

Hello,

Thank you for this excellent work!
I’ve recently been experimenting with the Aurora model to forecast 2-meter temperature (t2m) and total precipitation (tp). As part of this, I’m trying to better understand the proper methodology for fine-tuning the Aurora model.

I’ve attached the code I wrote for this task. Could you please take a look and let me know if my approach is correct? In particular, I applied normalization to the t2m variable, but I’d like to confirm whether this is the right way to handle it.
Additionally, could you clarify which input data variables are most relevant for forecasting t2m?
Any feedback or guidance you can provide on improving the fine-tuning process would be greatly appreciated.
Thank you very much for your time and support.

aurora_finetune_t2m.py

from datetime import datetime
import xarray as xr
import torch
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from aurora import AuroraPretrained, Batch, Metadata, rollout
from aurora.normalisation import locations, scales

=============================

1) Load datasets

=============================

static_ds = xr.open_dataset("./era5/static.nc", engine="netcdf4")
surf_ds = xr.open_dataset("./era5/2024_09_surface-level.nc", engine="netcdf4") # expects vars: t2m, u10, v10, msl
atm_ds = xr.open_dataset("./era5/2024_09_atmos.nc", engine="netcdf4") # expects vars: z,u,v,t,q at pressure_level

assert "t2m" in surf_ds and "u10" in surf_ds and "v10" in surf_ds and "msl" in surf_ds, "surface.nc missing required keys"
for k in ["z", "u", "v", "t", "q"]:
assert k in atm_ds, f"2024_09_atmos.nc missing {k}"

lat = torch.tensor(surf_ds.latitude.values, dtype=torch.float32)
lon = torch.tensor(surf_ds.longitude.values, dtype=torch.float32)
levels = tuple(int(l) for l in atm_ds.pressure_level.values)

train_steps = len(surf_ds.valid_time) - 1
print("Train steps:", train_steps)

=============================

2) Normalization stats

(update locations/scales for all used vars)

=============================

def update_stats(name, arr):
locations[name] = float(arr.mean())
std = float(arr.std())
scales[name] = std if std > 1e-6 else 1.0

update_stats("2t", surf_ds["t2m"].values)
update_stats("10u", surf_ds["u10"].values)
update_stats("10v", surf_ds["v10"].values)
update_stats("msl", surf_ds["msl"].values)
for k in ["z","u","v","t","q"]:
update_stats(k, atm_ds[k].values)

def norm(x, name): return (x - locations[name]) / (scales[name] + 1e-6)
def denorm(x, name): return x * (scales[name] + 1e-6) + locations[name]

=============================

3) Dataset

=============================

class T2MDataset(Dataset):
def init(self, static, surf, atm, lat, lon, levels):
self.static = static
self.surf = surf
self.atm = atm
self.lat = lat
self.lon = lon
self.levels = levels
self.steps = len(self.surf.valid_time) - 1

    self.static_tensors = {
        "z":   torch.tensor(self.static["z"].values[0], dtype=torch.float32),
        "slt": torch.tensor(self.static["slt"].values[0], dtype=torch.float32),
        "lsm": torch.tensor(self.static["lsm"].values[0], dtype=torch.float32),
    }

def __len__(self): return self.steps

def __getitem__(self, idx):
    surf_vars = {
        "2t":  norm(torch.tensor(self.surf["t2m"].values[idx], dtype=torch.float32), "2t").unsqueeze(0).unsqueeze(0),
        "10u": norm(torch.tensor(self.surf["u10"].values[idx], dtype=torch.float32), "10u").unsqueeze(0).unsqueeze(0),
        "10v": norm(torch.tensor(self.surf["v10"].values[idx], dtype=torch.float32), "10v").unsqueeze(0).unsqueeze(0),
        "msl": norm(torch.tensor(self.surf["msl"].values[idx], dtype=torch.float32), "msl").unsqueeze(0).unsqueeze(0),
    }

    atmos_vars = {
        k: norm(torch.tensor(self.atm[k].values[idx], dtype=torch.float32), k).unsqueeze(0).unsqueeze(0)
        for k in ["z","u","v","t","q"]
    }
    
    static_vars = {k: v.clone() for k, v in self.static_tensors.items()}

    meta = Metadata(
        lat=self.lat,
        lon=self.lon,
        time=(pd.to_datetime(str(self.surf.valid_time.values[idx])),),
        atmos_levels=self.levels
    )

    target = norm(torch.tensor(self.surf["t2m"].values[idx+1], dtype=torch.float32), "2t").unsqueeze(0).unsqueeze(0)
    
    return surf_vars, static_vars, atmos_vars, meta, target

dataset = T2MDataset(static_ds, surf_ds, atm_ds, lat, lon, levels)

Adjust loader workers/pinning to your machine

loader = DataLoader(
dataset, batch_size=1, shuffle=True, num_workers=2, pin_memory=True
)

=============================

4) Model, Optimizer, AMP

=============================

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

model = AuroraPretrained(
surf_vars=("2t","10u","10v","msl"),
static_vars=("lsm","z","slt"),
atmos_vars=("z","u","v","t","q"),
bf16_mode=True
).to(device)
model.load_checkpoint(strict=False)

optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4, weight_decay=1e-4)
criterion = torch.nn.HuberLoss(delta=1.0) # robust vs spikes
scaler = torch.cuda.amp.GradScaler(enabled=(device.type=="cuda"))

=============================

5) Train

=============================

epochs = 20
max_grad_norm = 1.0

for epoch in range(epochs):
model.train()
total = 0.0

for surf_vars, static_vars, atmos_vars, meta, target in loader:
    surf_vars = {k: v.to(device, non_blocking=True) for k, v in surf_vars.items()}
    static_vars = {k: v.to(device, non_blocking=True) for k, v in static_vars.items()}
    atmos_vars = {k: v.to(device, non_blocking=True) for k, v in atmos_vars.items()}
    meta = Metadata(
        lat=meta.lat.to(device, non_blocking=True),
        lon=meta.lon.to(device, non_blocking=True),
        time=meta.time,
        atmos_levels=meta.atmos_levels,
    )
    target = target.to(device, non_blocking=True)

    batch = Batch(surf_vars=surf_vars, static_vars=static_vars, atmos_vars=atmos_vars, metadata=meta)
    optimizer.zero_grad(set_to_none=True)
    with torch.cuda.amp.autocast(enabled=(device.type=="cuda")):
        pred = model(batch)  # pred.surf_vars["2t"] in normalized space
        loss = criterion(pred.surf_vars["2t"], target)

    scaler.scale(loss).backward()
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)
    scaler.step(optimizer)
    scaler.update()

    total += loss.item()

avg = total / len(loader)
print(f"Epoch {epoch+1}/{epochs} - Loss: {avg:.6f}")

torch.save({
    "epoch": epoch + 1,
    "model_state": model.state_dict(),
    "optim_state": optimizer.state_dict(),
    "norm_stats": {"locations": dict(locations), "scales": dict(scales)},
}, "aurora_t2m_finetuned.pt")

print("✅ Training complete. Checkpoint saved: aurora_t2m_finetuned.pt")

Start by reviewing the attached aurora_finetune_t2m.py script alongside the AuroraPretrained, Batch, Metadata, and aurora.normalisation entry points it uses. Check the normalization statistics and selected input variables, then document whether the fine-tuning approach is valid and what a correct, reproducible workflow should produce.