Update SpectralGPT model package

conf/config.yaml

@@ -1,28 +1,42 @@
 model:
-  image_size: 24
   in_channels: 12
   patch_size: 8
   spectral_patch_size: 3
-  embed_dim: 48
-  encoder_depth: 2
+  embed_dim: 32
+  encoder_depth: 1
   encoder_heads: 4
   decoder_dim: 32
   decoder_depth: 1
   decoder_heads: 4
   mask_ratio: 0.90
-  spectral_loss_weight: 1.0
+  spectral_angle_weight: 0.1
+  spectral_gradient_weight: 0.1
+stages:
+  - name: stage1
+    dataset: fMoW-Sentinel
+    image_size: 96
+    train_path: ./data/stage1_train.npz
+    train_samples: 2
+    epochs: 1
+  - name: stage2
+    dataset: BigEarthNet
+    image_size: 128
+    train_path: ./data/stage2_train.npz
+    train_samples: 2
+    epochs: 1
 data:
-  path: ./data/fake_spectralgpt.npz
-  samples: 8
-  protocol: synthetic_sentinel2_npz
+  test_path: ./data/stage2_test.npz
+  test_samples: 1
+  protocol: spectralgpt_progressive_s2_v2
 training:
-  epochs: 2
-  batch_size: 2
+  batch_size: 1
   learning_rate: 0.0001
   weight_decay: 0.05
   save_dir: ./result/checkpoints
-  checkpoint: ./result/checkpoints/best.pth
+  checkpoint: ./result/checkpoints/final.pth
+  metrics: ./result/training/metrics.json
+  amp: true
 runtime:
-  device: auto
+  device: cpu
   seed: 42
   output_dir: ./result/output

config.json

@@ -5,42 +5,62 @@
     "SpectralGPT"
   ],
   "framework": "PyTorch",
-  "domain": "earth-science",
+  "domain": "earth-observation",
   "task": "remote-sensing-masked-image-modeling",
   "implementation": {
     "entry_point": "model/spectralgpt.py",
-    "scope": "compact 12-band Sentinel-2 masked autoencoder reproduction"
+    "scope": "compact spatial-spectral masked autoencoder with progressive two-stage pretraining"
   },
   "architecture": {
-    "family": "spectral-spatial masked autoencoder",
-    "input_format": "NCHW multispectral images",
+    "family": "3D spatial-spectral masked autoencoder",
     "input_channels": 12,
-    "image_size": 24,
+    "stage_image_shapes": [
+      [
+        96,
+        96
+      ],
+      [
+        128,
+        128
+      ]
+    ],
     "patch_size": 8,
     "spectral_patch_size": 3,
-    "embed_dim": 48,
-    "encoder_depth": 2,
+    "embed_dim": 32,
+    "encoder_depth": 1,
+    "encoder_heads": 4,
     "decoder_dim": 32,
     "decoder_depth": 1,
+    "decoder_heads": 4,
     "mask_ratio": 0.9,
-    "training_objective": "masked reconstruction with spectral loss"
+    "spectral_angle_weight": 0.1,
+    "spectral_gradient_weight": 0.1
   },
   "data": {
     "datasets": [
-      "fMoW-S2",
-      "BigEarthNet-S2"
+      "fMoW-Sentinel",
+      "BigEarthNet"
+    ],
+    "protocol": "spectralgpt_progressive_s2_v2",
+    "input_format": "NCHW_NPZ",
+    "stages": [
+      {
+        "name": "stage1",
+        "dataset": "fMoW-Sentinel",
+        "image_size": 96,
+        "train_samples": 2
+      },
+      {
+        "name": "stage2",
+        "dataset": "BigEarthNet",
+        "image_size": 128,
+        "train_samples": 2
+      }
     ],
-    "official_paper": "https://arxiv.org/abs/2311.07113",
-    "protocol": "synthetic_sentinel2_npz",
-    "synthetic_samples": 8
+    "test_samples": 1
   },
-  "metrics": [
-    "mse",
-    "mae",
-    "psnr_db",
-    "per_band_rmse"
-  ],
   "configuration_sources": [
+    "configuration.json",
     "conf/config.yaml",
     "model/spectralgpt.py"
   ]

configuration.json

@@ -1,9 +1,10 @@
 {
   "framework": "PyTorch",
+  "license": "GPL-3.0",
   "task": "remote_sensing_masked_image_modeling",
   "model": "SpectralGPT",
   "input_format": "NCHW_NPZ",
-  "protocol": "synthetic_sentinel2_npz",
+  "protocol": "spectralgpt_progressive_s2_v2",
   "default_config": "conf/config.yaml",
   "train": "scripts/train.py",
   "inference": "scripts/inference.py",

model/spectralgpt.py

@@ -8,7 +8,8 @@ class SpectralGPT(nn.Module):
     def __init__(self, image_size=24, in_channels=12, patch_size=8,
                  spectral_patch_size=3, embed_dim=48, encoder_depth=2,
                  encoder_heads=4, decoder_dim=32, decoder_depth=1,
-                 decoder_heads=4, mask_ratio=0.9, spectral_loss_weight=1.0):
+                  decoder_heads=4, mask_ratio=0.9, spectral_angle_weight=0.1,
+                  spectral_gradient_weight=0.1):
         super().__init__()
         if image_size % patch_size or in_channels % spectral_patch_size:
             raise ValueError("Image and spectral dimensions must be divisible by token sizes")
@@ -21,7 +22,8 @@ class SpectralGPT(nn.Module):
         self.num_tokens = self.spatial_tokens * self.spectral_tokens
         self.token_pixels = patch_size * patch_size * spectral_patch_size
         self.mask_ratio = mask_ratio
-        self.spectral_loss_weight = spectral_loss_weight
+        self.spectral_angle_weight = spectral_angle_weight
+        self.spectral_gradient_weight = spectral_gradient_weight
 
         self.patch_embed = nn.Conv3d(
             1, embed_dim,
@@ -96,17 +98,31 @@ class SpectralGPT(nn.Module):
         prediction = self.decoder_pred(self.decoder_norm(self.decoder(full + self.decoder_pos(positions))))
         target = self.patchify(images)
         token_error = (prediction - target).pow(2).mean(-1)
-        token_loss = (token_error * mask).sum() / mask.sum().clamp_min(1)
-        n = images.shape[0]
-        predicted_grid = prediction.reshape(n, self.spectral_tokens, self.spatial_tokens, -1)
-        target_grid = target.reshape(n, self.spectral_tokens, self.spatial_tokens, -1)
-        spectral_loss = (predicted_grid[:, 1:] - target_grid[:, 1:]).pow(2).mean()
-        loss = token_loss + self.spectral_loss_weight * spectral_loss
+        masked_mse = (token_error * mask).sum() / mask.sum().clamp_min(1)
+        mask_image = self.unpatchify(mask.unsqueeze(-1).expand(-1, -1, self.token_pixels))
+        predicted_image = self.unpatchify(prediction)
+        completed = images * (1.0 - mask_image) + predicted_image * mask_image
+        spectral_mask = mask_image.any(dim=1)
+        completed_norm = completed.norm(dim=1)
+        target_norm = images.norm(dim=1)
+        valid_sam = spectral_mask & (completed_norm > 1e-6) & (target_norm > 1e-6)
+        cosine = (completed * images).sum(dim=1) / (completed_norm * target_norm).clamp_min(1e-6)
+        angles = torch.acos(cosine.clamp(-1.0, 1.0))
+        spectral_angle = (angles * valid_sam).sum() / valid_sam.sum().clamp_min(1)
+        completed_gradient = completed[:, 1:] - completed[:, :-1]
+        target_gradient = images[:, 1:] - images[:, :-1]
+        gradient_mask = torch.maximum(mask_image[:, 1:], mask_image[:, :-1]).bool()
+        spectral_gradient = ((completed_gradient - target_gradient).abs() * gradient_mask).sum() / gradient_mask.sum().clamp_min(1)
+        loss = (masked_mse + self.spectral_angle_weight * spectral_angle
+                + self.spectral_gradient_weight * spectral_gradient)
         return {
             "loss": loss,
-            "token_loss": token_loss,
-            "spectral_loss": spectral_loss,
+            "masked_mse": masked_mse,
+            "spectral_angle": spectral_angle,
+            "spectral_gradient": spectral_gradient,
             "prediction": prediction,
             "mask": mask,
-            "reconstruction": self.unpatchify(prediction),
+            "mask_image": mask_image,
+            "prediction_image": predicted_image,
+            "reconstruction": completed,
         }

scripts/fake_data.py

@@ -5,15 +5,16 @@ import numpy as np
 import yaml
 
 
-def real_images(directory, size, samples):
+def real_images(paths, size, scale_factor):
     import tifffile
 
-    paths = sorted(Path(directory).rglob("*.tif"))[:samples]
     if not paths:
-        raise FileNotFoundError(f"No TIFF files found under {directory}")
+        raise FileNotFoundError("No TIFF files supplied")
     images = []
+    scale_factors = []
     for path in paths:
-        image = tifffile.imread(path).astype(np.float32)
+        image = tifffile.imread(path)
+        original_dtype = image.dtype
         if image.ndim != 3:
             raise ValueError(f"Expected a 13-band TIFF, got {image.shape} from {path}")
         if image.shape[0] == 13:
@@ -24,43 +25,86 @@ def real_images(directory, size, samples):
         y = np.linspace(0, image.shape[0] - 1, size).round().astype(int)
         x = np.linspace(0, image.shape[1] - 1, size).round().astype(int)
         image = image[y][:, x]
-        images.append(np.clip(image / 10000.0, 0, 1).transpose(2, 0, 1))
+        factor = scale_factor
+        if factor is None:
+            minimum = float(np.nanmin(image))
+            maximum = float(np.nanmax(image))
+            already_normalized = minimum >= 0.0 and maximum <= 1.0
+            factor = 10000.0 if not already_normalized and (
+                np.issubdtype(original_dtype, np.integer) or maximum > 1.0
+            ) else 1.0
+        images.append(np.clip(image.astype(np.float32) / factor, 0, 1).transpose(2, 0, 1))
+        scale_factors.append(factor)
+    return np.asarray(images, dtype=np.float32), np.asarray(scale_factors, dtype=np.float32)
+
+
+def synthetic_images(count, size, seed):
+    rng = np.random.default_rng(seed)
+    y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
+    images = []
+    for index in range(count):
+        phase = rng.uniform(0, 2 * np.pi)
+        bands = []
+        for band in range(12):
+            pattern = 0.45 + 0.22 * np.sin((band + 1) * x + phase)
+            pattern += 0.18 * np.cos((band / 3 + 1) * y - phase)
+            pattern += rng.normal(0, 0.025, (size, size))
+            bands.append(np.clip(pattern, 0, 1))
+        images.append(bands)
     return np.asarray(images, dtype=np.float32)
 
 
+def save_npz(output, images, source, protocol, normalization, scale_factors, stage):
+    band_order = np.asarray(["B1", "B2", "B3", "B4", "B5", "B6", "B7", "B8", "B8A", "B9", "B11", "B12"])
+    output.parent.mkdir(parents=True, exist_ok=True)
+    np.savez_compressed(output, images=images, data_source=np.asarray(source),
+                        protocol=np.asarray(protocol), band_order=band_order,
+                        normalization=np.asarray(normalization), scale_factors=scale_factors,
+                        stage=np.asarray(stage))
+    print(f"saved: {output} shape={images.shape} stage={stage} data_source={source}")
+
+
 def main():
     parser = argparse.ArgumentParser(description="Generate compact 12-band spectral data")
     parser.add_argument("--config", default="conf/config.yaml")
     parser.add_argument("--real-dir", help="Convert official 13-band Sentinel-2 TIFF files")
+    parser.add_argument("--scale-factor", default="auto",
+                        help="TIFF divisor, or 'auto' (10000 for integer/range > 1; otherwise 1)")
+    parser.add_argument("--stage", choices=("stage1", "stage2"), default="stage2",
+                        help="Target stage for real TIFF conversion")
     args = parser.parse_args()
     with open(args.config, encoding="utf-8") as handle:
         config = yaml.safe_load(handle)
-    size = config["model"]["image_size"]
-    samples = config["data"]["samples"]
+    if args.scale_factor == "auto":
+        scale_factor = None
+    else:
+        scale_factor = float(args.scale_factor)
+        if not np.isfinite(scale_factor) or scale_factor <= 0:
+            raise ValueError("--scale-factor must be a positive finite number or 'auto'")
     if args.real_dir:
-        images = real_images(args.real_dir, size, samples)
-        source = "real"
-        protocol = "real_sentinel2_npz"
+        stage = next(item for item in config["stages"] if item["name"] == args.stage)
+        count = stage["train_samples"] + (config["data"]["test_samples"] if args.stage == "stage2" else 0)
+        paths = sorted(Path(args.real_dir).rglob("*.tif"))
+        if len(paths) < count:
+            raise ValueError(f"Need at least {count} TIFF files, found {len(paths)}")
+        images, scale_factors = real_images(paths[:count], stage["image_size"], scale_factor)
+        normalization = "divide_by_scale_factor_then_clip_0_1"
+        save_npz(Path(stage["train_path"]), images[:stage["train_samples"]], "real",
+                 config["data"]["protocol"], normalization,
+                 scale_factors[:stage["train_samples"]], args.stage)
+        if args.stage == "stage2":
+            save_npz(Path(config["data"]["test_path"]), images[stage["train_samples"]:], "real",
+                     config["data"]["protocol"], normalization,
+                     scale_factors[stage["train_samples"]:], "stage2")
     else:
-        rng = np.random.default_rng(config["runtime"]["seed"])
-        y, x = np.mgrid[0:size, 0:size].astype(np.float32) / max(size - 1, 1)
-        images = []
-        for index in range(samples):
-            phase = rng.uniform(0, 2 * np.pi)
-            bands = []
-            for band in range(12):
-                pattern = 0.45 + 0.22 * np.sin((band + 1) * x + phase)
-                pattern += 0.18 * np.cos((band / 3 + 1) * y - phase)
-                pattern += rng.normal(0, 0.025, (size, size))
-                bands.append(np.clip(pattern, 0, 1))
-            images.append(bands)
-        images = np.asarray(images, dtype=np.float32)
-        source = "synthetic"
-        protocol = config["data"]["protocol"]
-    output = Path(config["data"]["path"])
-    output.parent.mkdir(parents=True, exist_ok=True)
-    np.savez_compressed(output, images=images, data_source=np.asarray(source), protocol=np.asarray(protocol))
-    print(f"saved: {output} shape={images.shape} data_source={source} protocol={protocol}")
+        for index, stage in enumerate(config["stages"]):
+            images = synthetic_images(stage["train_samples"], stage["image_size"], config["runtime"]["seed"] + index)
+            save_npz(Path(stage["train_path"]), images, "synthetic", config["data"]["protocol"],
+                     "already_0_1", np.ones(len(images), np.float32), stage["name"])
+        test_size = config["stages"][-1]["image_size"]
+        images = synthetic_images(config["data"]["test_samples"], test_size, config["runtime"]["seed"] + 2)
+        save_npz(Path(config["data"]["test_path"]), images, "synthetic", config["data"]["protocol"],
+                 "already_0_1", np.ones(len(images), np.float32), "stage2")
 
 
 if __name__ == "__main__":

scripts/inference.py

@@ -14,40 +14,61 @@ def main():
     parser = argparse.ArgumentParser(description="Run SpectralGPT reconstruction")
     parser.add_argument("--config", default="conf/config.yaml")
     parser.add_argument("--checkpoint")
+    parser.add_argument("--batch-size", type=int)
     args = parser.parse_args()
     with open(args.config, encoding="utf-8") as handle:
         config = yaml.safe_load(handle)
     requested = config["runtime"]["device"]
     device = torch.device("cuda" if torch.cuda.is_available() and requested != "cpu" else "cpu")
     torch.manual_seed(config["runtime"]["seed"])
-    model = SpectralGPT(**config["model"]).to(device)
+    stage = config["stages"][-1]
+    model = SpectralGPT(image_size=stage["image_size"], **config["model"]).to(device)
     checkpoint_path = args.checkpoint or config["training"]["checkpoint"]
     if not Path(checkpoint_path).exists():
         raise FileNotFoundError(
             f"Missing checkpoint: {checkpoint_path}. Run `python scripts/train.py` first."
         )
     checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    if checkpoint.get("stage") != stage["name"] or checkpoint.get("image_size") != stage["image_size"]:
+        raise ValueError("Checkpoint is not the configured final stage2 checkpoint")
     model.load_state_dict(checkpoint["model"])
     model.eval()
-    data_path = Path(config["data"]["path"])
+    data_path = Path(config["data"]["test_path"])
     if not data_path.exists():
         raise FileNotFoundError(
             f"Missing inference data: {data_path}. Run `python scripts/fake_data.py` first."
         )
     with np.load(data_path) as data:
-        images = torch.from_numpy(data["images"]).to(device)
+        images = data["images"].copy()
         data_source = str(data["data_source"]) if "data_source" in data.files else "unknown"
         protocol = str(data["protocol"]) if "protocol" in data.files else "unknown"
-    with torch.no_grad():
-        output = model(images)
+        normalization = str(data["normalization"]) if "normalization" in data.files else "unknown"
+        scale_factors = data["scale_factors"].copy() if "scale_factors" in data.files else np.ones(len(images), np.float32)
+        stored_stage = str(data["stage"]) if "stage" in data.files else "unknown"
+    expected = (config["model"]["in_channels"], stage["image_size"], stage["image_size"])
+    if images.dtype != np.float32 or images.ndim != 4 or tuple(images.shape[1:]) != expected:
+        raise ValueError(f"Expected float32 stage2 test [N,{','.join(map(str, expected))}], got {images.dtype} {images.shape}")
+    if stored_stage != stage["name"]:
+        raise ValueError(f"Expected test stage {stage['name']}, got {stored_stage}")
+    batch_size = args.batch_size or config["training"]["batch_size"]
+    collected = {name: [] for name in ("reconstruction", "prediction_image", "mask", "mask_image")}
+    with torch.inference_mode():
+        for start in range(0, len(images), batch_size):
+            batch = torch.from_numpy(images[start:start + batch_size]).to(device)
+            output = model(batch)
+            for name in collected:
+                collected[name].append(output[name].cpu().numpy())
     output_dir = Path(config["runtime"]["output_dir"])
     output_dir.mkdir(parents=True, exist_ok=True)
     np.savez_compressed(output_dir / "reconstruction.npz",
-                        inputs=images.cpu().numpy(),
-                        reconstructions=output["reconstruction"].cpu().numpy(),
-                        masks=output["mask"].cpu().numpy(),
-                        data_source=np.asarray(data_source),
-                        protocol=np.asarray(protocol))
+                         inputs=images,
+                         reconstructions=np.concatenate(collected["reconstruction"]),
+                         predictions=np.concatenate(collected["prediction_image"]),
+                         masks=np.concatenate(collected["mask"]),
+                         mask_images=np.concatenate(collected["mask_image"]),
+                         data_source=np.asarray(data_source),
+                         protocol=np.asarray(protocol), normalization=np.asarray(normalization),
+                         scale_factors=scale_factors, stage=np.asarray(stored_stage))
     print(
         f"saved: {output_dir / 'reconstruction.npz'} "
         f"data_source={data_source} protocol={protocol}"

scripts/result.py

@@ -4,6 +4,7 @@ from pathlib import Path
 
 import numpy as np
 import yaml
+import matplotlib.pyplot as plt
 
 
 def rgb(image):
@@ -27,24 +28,41 @@ def main():
     with np.load(reconstruction_path) as data:
         inputs = data["inputs"]
         reconstructions = data["reconstructions"]
+        predictions = data["predictions"]
+        mask_images = data["mask_images"]
         data_source = str(data["data_source"]) if "data_source" in data.files else "unknown"
         protocol = str(data["protocol"]) if "protocol" in data.files else "unknown"
-    mse = float(np.mean((inputs - reconstructions) ** 2))
-    mae = float(np.mean(np.abs(inputs - reconstructions)))
+        normalization = str(data["normalization"]) if "normalization" in data.files else "unknown"
+    denominator = max(float(mask_images.sum()), 1.0)
+    mse = float((((inputs - predictions) ** 2) * mask_images).sum() / denominator)
+    mae = float((np.abs(inputs - predictions) * mask_images).sum() / denominator)
     psnr = float(-10 * np.log10(max(mse, 1e-12)))
-    spectral_rmse = np.sqrt(np.mean((inputs - reconstructions) ** 2, axis=(0, 2, 3)))
-    metrics = {"mse": mse, "mae": mae, "psnr_db": psnr,
+    per_band_denominator = np.maximum(mask_images.sum(axis=(0, 2, 3)), 1)
+    spectral_rmse = np.sqrt((((inputs - predictions) ** 2) * mask_images).sum(axis=(0, 2, 3)) / per_band_denominator)
+    dot = (inputs * reconstructions).sum(axis=1)
+    norms = np.linalg.norm(inputs, axis=1) * np.linalg.norm(reconstructions, axis=1)
+    pixel_mask = (mask_images > 0).any(axis=1) & (norms > 1e-8)
+    sam = np.arccos(np.clip(dot / np.maximum(norms, 1e-8), -1, 1))
+    metrics = {"masked_mse": mse, "masked_mae": mae, "masked_psnr_db": psnr,
+               "masked_spectral_angle_deg": float(np.degrees(sam[pixel_mask]).mean()),
                "data_source": data_source, "protocol": protocol,
+               "normalization": normalization,
                "per_band_rmse": spectral_rmse.tolist()}
     with open(output_dir / "metrics.json", "w", encoding="utf-8") as handle:
         json.dump(metrics, handle, indent=2)
-    panel = np.concatenate([rgb(inputs[0]), rgb(reconstructions[0])], axis=1)
-    with open(output_dir / "reconstruction.ppm", "wb") as handle:
-        handle.write(f"P6\n{panel.shape[1]} {panel.shape[0]}\n255\n".encode("ascii"))
-        handle.write(panel.tobytes())
+    masked = inputs[0] * (1.0 - mask_images[0])
+    figure, axes = plt.subplots(1, 4, figsize=(13, 3.5))
+    for axis, image, title in zip(axes, [inputs[0], masked, predictions[0], reconstructions[0]],
+                                  ["Input", "Visible tokens", "MAE prediction", "Composite"]):
+        axis.imshow(rgb(image))
+        axis.set_title(title)
+        axis.axis("off")
+    figure.tight_layout()
+    figure.savefig(output_dir / "reconstruction.png", dpi=140)
+    plt.close(figure)
     print(json.dumps(metrics, indent=2))
     print(f"saved: {output_dir / 'metrics.json'}")
-    print(f"saved: {output_dir / 'reconstruction.ppm'}")
+    print(f"saved: {output_dir / 'reconstruction.png'}")
 
 
 if __name__ == "__main__":

scripts/train.py

@@ -1,10 +1,13 @@
 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
@@ -14,13 +17,20 @@ from model.spectralgpt import SpectralGPT
 
 
 class SpectralDataset(Dataset):
-    def __init__(self, path):
+    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 the images array")
+                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)
@@ -29,85 +39,105 @@ class SpectralDataset(Dataset):
         return torch.from_numpy(self.images[index])
 
 
-def load_config(path):
-    with open(path, encoding="utf-8") as handle:
-        return yaml.safe_load(handle)
-
-
-def build_model(config):
-    return SpectralGPT(**config["model"])
+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="Train compact SpectralGPT")
+    parser = argparse.ArgumentParser(description="Progressive two-stage SpectralGPT training")
     parser.add_argument("--config", default="conf/config.yaml")
-    parser.add_argument("--data")
-    parser.add_argument("--epochs", type=int)
     args = parser.parse_args()
-    config = load_config(args.config)
+    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"))
-    if distributed:
-        torch.distributed.init_process_group("nccl" if torch.cuda.is_available() else "gloo")
     requested = config["runtime"]["device"]
-    use_accelerator = torch.cuda.is_available() and requested != "cpu"
-    device = torch.device(f"cuda:{local_rank}" if use_accelerator else "cpu")
-    torch.manual_seed(config["runtime"]["seed"] + local_rank)
-    data_path = Path(args.data or config["data"]["path"])
-    if not data_path.exists():
-        raise FileNotFoundError(
-            f"Missing training data: {data_path}. "
-            "Run `python scripts/fake_data.py` for a synthetic connectivity test."
-        )
-    dataset = SpectralDataset(data_path)
-    if dataset.images.ndim != 4 or tuple(dataset.images.shape[1:]) != (
-        config["model"]["in_channels"], config["model"]["image_size"], config["model"]["image_size"]
-    ):
-        raise ValueError(
-            f"Expected images shaped [N,{config['model']['in_channels']},"
-            f"{config['model']['image_size']},{config['model']['image_size']}], "
-            f"got {dataset.images.shape}"
-        )
-    if local_rank == 0:
-        print(
-            f"data_source={dataset.data_source} protocol={dataset.protocol} "
-            f"samples={len(dataset)}"
-        )
-    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 = build_model(config).to(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:
-        model = DistributedDataParallel(model, device_ids=[local_rank] if use_accelerator else None)
-    optimizer = torch.optim.AdamW(model.parameters(), lr=config["training"]["learning_rate"],
-                                  weight_decay=config["training"]["weight_decay"], betas=(0.9, 0.95))
+        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"])
-    best = float("inf")
-    for epoch in range(args.epochs or config["training"]["epochs"]):
-        if sampler is not None:
-            sampler.set_epoch(epoch)
-        model.train()
-        losses = []
-        for images in loader:
-            output = model(images.to(device))
-            optimizer.zero_grad()
-            output["loss"].backward()
-            optimizer.step()
-            losses.append(output["loss"].item())
-        mean_loss = float(np.mean(losses))
-        if local_rank == 0:
-            print(f"epoch={epoch + 1} reconstruction_loss={mean_loss:.6f}")
+    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)
-            state = model.module.state_dict() if distributed else model.state_dict()
-            checkpoint = {"model": state, "epoch": epoch + 1, "loss": mean_loss,
+            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,
-                          "format": "spectralgpt-compact-v1"}
-            torch.save(checkpoint, save_dir / "last.pth")
-            if mean_loss < best:
-                best = mean_loss
-                torch.save(checkpoint, save_dir / "best.pth")
-    if local_rank == 0:
-        print(f"saved: {save_dir / 'best.pth'}")
+                          "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()