import argparse import json import math import os from pathlib import Path import sys import numpy as np import torch import torch.nn.functional as F import yaml from torch.nn.parallel import DistributedDataParallel from torch.utils.data import DataLoader, Dataset, DistributedSampler sys.path.insert(0, str(Path(__file__).resolve().parents[1])) from model.spectralgpt import SpectralGPT class SpectralDataset(Dataset): def __init__(self, path, image_size, stage): with np.load(path) as data: if "images" not in data.files: raise ValueError(f"Dataset {path} is missing images") self.images = data["images"].copy() self.data_source = str(data["data_source"]) if "data_source" in data.files else "unknown" self.protocol = str(data["protocol"]) if "protocol" in data.files else "unknown" self.normalization = str(data["normalization"]) if "normalization" in data.files else "unknown" stored_stage = str(data["stage"]) if "stage" in data.files else "unknown" expected = (12, image_size, image_size) if self.images.dtype != np.float32 or self.images.ndim != 4 or tuple(self.images.shape[1:]) != expected: raise ValueError(f"Expected float32 [N,{','.join(map(str, expected))}], got {self.images.dtype} {self.images.shape}") if stored_stage != stage: raise ValueError(f"Expected stage metadata {stage}, got {stored_stage}") def __len__(self): return len(self.images) def __getitem__(self, index): return torch.from_numpy(self.images[index]) def resize_spatial_position(state, old_size, new_size, patch_size): if old_size == new_size: return state key = "spatial_pos" position = state[key] old_grid, new_grid = old_size // patch_size, new_size // patch_size if position.shape[1] != old_grid * old_grid: raise ValueError("Checkpoint spatial position shape does not match previous stage") position = position.reshape(1, old_grid, old_grid, -1).permute(0, 3, 1, 2) state[key] = F.interpolate(position, size=(new_grid, new_grid), mode="bicubic", align_corners=False).permute(0, 2, 3, 1).reshape(1, new_grid * new_grid, -1) return state def main(): parser = argparse.ArgumentParser(description="Progressive two-stage SpectralGPT training") parser.add_argument("--config", default="conf/config.yaml") args = parser.parse_args() with open(args.config, encoding="utf-8") as handle: config = yaml.safe_load(handle) distributed = int(os.environ.get("WORLD_SIZE", "1")) > 1 rank = int(os.environ.get("RANK", "0")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) requested = config["runtime"]["device"] device = torch.device(f"cuda:{local_rank}" if torch.cuda.is_available() and requested != "cpu" else "cpu") if device.type == "cuda": torch.cuda.set_device(device) if distributed: torch.distributed.init_process_group("nccl" if device.type == "cuda" else "gloo") torch.manual_seed(config["runtime"]["seed"] + rank) amp_enabled = bool(config["training"].get("amp", True) and device.type == "cuda") save_dir = Path(config["training"]["save_dir"]) history = [] previous_state = None previous_size = None for stage in config["stages"]: path = Path(stage["train_path"]) if not path.exists(): raise FileNotFoundError(f"Missing {stage['name']} data: {path}. Run scripts/fake_data.py") dataset = SpectralDataset(path, stage["image_size"], stage["name"]) sampler = DistributedSampler(dataset, shuffle=True) if distributed else None loader = DataLoader(dataset, batch_size=config["training"]["batch_size"], sampler=sampler, shuffle=sampler is None) model = SpectralGPT(image_size=stage["image_size"], **config["model"]) if previous_state is not None: model.load_state_dict(resize_spatial_position(previous_state, previous_size, stage["image_size"], config["model"]["patch_size"])) model = model.to(device) if distributed: model = DistributedDataParallel(model, device_ids=[local_rank] if device.type == "cuda" else None) optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"], weight_decay=config["training"]["weight_decay"], betas=(0.9, 0.95)) scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled) for epoch in range(stage["epochs"]): if sampler is not None: sampler.set_epoch(epoch) model.train() totals = torch.zeros(5, dtype=torch.float64, device=device) for images in loader: images = images.to(device) optimizer.zero_grad(set_to_none=True) with torch.autocast(device_type=device.type, dtype=torch.float16, enabled=amp_enabled): output = model(images) scaler.scale(output["loss"]).backward() scaler.step(optimizer) scaler.update() count = images.shape[0] totals += torch.tensor([output[name].item() * count for name in ("loss", "masked_mse", "spectral_angle", "spectral_gradient")] + [count], dtype=torch.float64, device=device) if distributed: torch.distributed.all_reduce(totals) values = (totals[:4] / totals[4]).tolist() record = {"stage": stage["name"], "dataset": stage["dataset"], "image_size": stage["image_size"], "patch_size": config["model"]["patch_size"], "epoch": epoch + 1, **dict(zip(("loss", "masked_mse", "spectral_angle", "spectral_gradient"), values))} history.append(record) if rank == 0: print(f"stage={stage['name']} epoch={epoch + 1} size={stage['image_size']} loss={values[0]:.6f}") base_model = model.module if distributed else model previous_state = {key: value.detach().cpu() for key, value in base_model.state_dict().items()} previous_size = stage["image_size"] if rank == 0: save_dir.mkdir(parents=True, exist_ok=True) checkpoint = {"model": previous_state, "config": config, "stage": stage["name"], "image_size": stage["image_size"], "stage_history": history, "data_source": dataset.data_source, "protocol": dataset.protocol, "normalization": dataset.normalization, "backward_completed": True, "format": "spectralgpt-progressive-v2"} torch.save(checkpoint, save_dir / f"{stage['name']}.pth") if stage is config["stages"][-1]: torch.save(checkpoint, Path(config["training"]["checkpoint"])) if rank == 0: metrics = Path(config["training"]["metrics"]) metrics.parent.mkdir(parents=True, exist_ok=True) metrics.write_text(json.dumps({"stage_history": history, "backward_completed": True, "amp_enabled": amp_enabled}, indent=2) + "\n", encoding="utf-8") print(f"saved: {config['training']['checkpoint']}") if distributed: torch.distributed.destroy_process_group() if __name__ == "__main__": main()