"""Train compact SkySense on NPZ multi-modal temporal samples.""" import importlib.util import json import os import random from pathlib import Path import numpy as np import torch import yaml from torch import distributed as dist from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, Dataset, DistributedSampler ROOT = Path(__file__).resolve().parents[1] def load_model_class(): spec = importlib.util.spec_from_file_location("skysense_model", ROOT / "model" / "skysense.py") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module.SkySense def load_config(): with (ROOT / "conf" / "config.yaml").open(encoding="utf-8") as handle: return yaml.safe_load(handle) class NPZDataset(Dataset): def __init__(self, path): archive = np.load(path) sample_keys = {"hr", "s2", "s1", "dates_hr", "dates_s2", "dates_s1", "region", "labels"} missing = sample_keys.difference(archive.files) if missing: raise ValueError(f"Dataset {path} is missing arrays: {sorted(missing)}") self.arrays = {key: archive[key] for key in sample_keys} self.data_source = str(archive["data_source"]) if "data_source" in archive.files else "unknown" self.protocol = str(archive["protocol"]) if "protocol" in archive.files else "unknown" def __len__(self): return len(self.arrays["hr"]) def __getitem__(self, index): return {key: torch.as_tensor(value[index]) for key, value in self.arrays.items()} def setup_device(config): world_size = int(os.environ.get("WORLD_SIZE", "1")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) use_accelerator = torch.cuda.is_available() and config["runtime"].get("device", "auto") != "cpu" if world_size > 1: backend = "nccl" if use_accelerator else "gloo" dist.init_process_group(backend=backend) if use_accelerator: torch.cuda.set_device(local_rank) return torch.device("cuda", local_rank), world_size, local_rank return torch.device("cpu"), world_size, local_rank def main(): config = load_config() seed = config["seed"] random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) device, world_size, local_rank = setup_device(config) amp_enabled = bool(config["train"].get("amp", True) and device.type == "cuda") dataset_path = ROOT / config["data"]["root"] / "train.npz" if not dataset_path.exists(): raise FileNotFoundError( f"Missing training data: {dataset_path.relative_to(ROOT)}. " "Run `python scripts/fake_data.py` for a synthetic connectivity test." ) dataset = NPZDataset(dataset_path) expected = { "hr": (config["data"]["hr_timesteps"], config["data"]["hr_channels"], config["data"]["hr_size"], config["data"]["hr_size"]), "s2": (config["data"]["s2_timesteps"], config["data"]["s2_channels"], config["data"]["s2_size"], config["data"]["s2_size"]), "s1": (config["data"]["s1_timesteps"], config["data"]["s1_channels"], config["data"]["s1_size"], config["data"]["s1_size"]), "labels": (config["data"]["hr_size"], config["data"]["hr_size"]), "dates_hr": (config["data"]["hr_timesteps"],), "dates_s2": (config["data"]["s2_timesteps"],), "dates_s1": (config["data"]["s1_timesteps"],), "region": (), } for key, shape in expected.items(): if tuple(dataset.arrays[key].shape[1:]) != shape: raise ValueError(f"Expected {key} shaped [N,{','.join(map(str, shape))}], got {dataset.arrays[key].shape}") if len(dataset.arrays[key]) != len(dataset): raise ValueError(f"Array {key} has {len(dataset.arrays[key])} samples, expected {len(dataset)}") for key in ("hr", "s2", "s1"): if not np.issubdtype(dataset.arrays[key].dtype, np.floating): raise TypeError(f"{key} must use a floating dtype, got {dataset.arrays[key].dtype}") for key in ("dates_hr", "dates_s2", "dates_s1", "region", "labels"): if dataset.arrays[key].dtype != np.int64: raise TypeError(f"{key} must use int64, got {dataset.arrays[key].dtype}") for key in ("dates_hr", "dates_s2", "dates_s1"): if np.any((dataset.arrays[key] < 0) | (dataset.arrays[key] > 364)): raise ValueError(f"{key} must contain day-of-year values in [0, 364]") if np.any((dataset.arrays["region"] < 0) | (dataset.arrays["region"] >= config["model"]["num_regions"])): raise ValueError(f"region IDs must be in [0, {config['model']['num_regions'] - 1}]") if np.any((dataset.arrays["labels"] < 0) | (dataset.arrays["labels"] >= config["data"]["num_classes"])): raise ValueError(f"labels must be in [0, {config['data']['num_classes'] - 1}]") if local_rank == 0: print( f"data_source={dataset.data_source} protocol={dataset.protocol} " f"samples={len(dataset)} hr={config['data']['hr_size']} s2={config['data']['s2_size']} s1={config['data']['s1_size']}" ) sampler = DistributedSampler(dataset, shuffle=True) if world_size > 1 else None loader = DataLoader( dataset, batch_size=config["train"]["batch_size"], shuffle=sampler is None, sampler=sampler, num_workers=config["train"]["num_workers"], ) SkySense = load_model_class() model = SkySense( **config["model"], hr_channels=config["data"]["hr_channels"], s2_channels=config["data"]["s2_channels"], s1_channels=config["data"]["s1_channels"], num_classes=config["data"]["num_classes"], ).to(device) if world_size > 1: model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) optimizer = torch.optim.AdamW(model.parameters(), lr=config["train"]["learning_rate"], weight_decay=config["train"]["weight_decay"]) scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) final_loss = float("nan") final_segmentation = float("nan") final_alignment = float("nan") for epoch in range(config["train"]["epochs"]): if sampler is not None: sampler.set_epoch(epoch) model.train() totals = torch.zeros(4, dtype=torch.float64, device=device) for batch in loader: batch = {key: value.to(device) for key, value in batch.items()} optimizer.zero_grad(set_to_none=True) with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=amp_enabled): output = model(batch["hr"], batch["s2"], batch["s1"], batch["dates_hr"], batch["dates_s2"], batch["dates_s1"], batch["region"]) base_model = model.module if hasattr(model, "module") else model segmentation = torch.nn.functional.cross_entropy(output["logits"], batch["labels"]) alignment = base_model.cross_modal_alignment_loss(output["features"]) loss = segmentation + config["train"]["alignment_weight"] * alignment scaler.scale(loss).backward() scaler.step(optimizer) scaler.update() batch_size = batch["hr"].shape[0] totals += torch.tensor([loss.item() * batch_size, segmentation.item() * batch_size, alignment.item() * batch_size, batch_size], dtype=torch.float64, device=device) if world_size > 1: dist.all_reduce(totals, op=dist.ReduceOp.SUM) final_loss = float((totals[0] / totals[3]).item()) final_segmentation = float((totals[1] / totals[3]).item()) final_alignment = float((totals[2] / totals[3]).item()) if local_rank == 0: print(f"epoch={epoch + 1} loss={final_loss:.6f}") if local_rank == 0: checkpoint = ROOT / config["paths"]["checkpoint"] checkpoint.parent.mkdir(parents=True, exist_ok=True) base_model = model.module if hasattr(model, "module") else model torch.save({"model": base_model.state_dict(), "config": config, "final_loss": final_loss}, checkpoint) metrics = ROOT / config["paths"]["training_metrics"] metrics.parent.mkdir(parents=True, exist_ok=True) metrics.write_text( json.dumps( { "final_loss": final_loss, "segmentation_loss": final_segmentation, "alignment_loss": final_alignment, "backward_completed": True, "amp_enabled": amp_enabled, "epochs": config["train"]["epochs"], "samples": len(dataset), "data_source": dataset.data_source, "protocol": dataset.protocol, }, indent=2, ) + "\n", encoding="utf-8", ) print(f"checkpoint={checkpoint.relative_to(ROOT)} final_loss={final_loss:.6f}") if world_size > 1: dist.destroy_process_group() if __name__ == "__main__": main()