Upload Scale-MAE model package

.gitattributes

@@ -1,35 +1,13 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
 *.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
 *.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text

conf/config.yaml

@@ -0,0 +1,44 @@
+seed: 42
+data:
+  root: data
+  protocol: synthetic_multiscale_geospatial_rgb
+  input_size: 16
+  target_size: 32
+  channels: 3
+  num_classes: 4
+  train_samples: 40
+  test_samples: 16
+  gsd_values: [0.5, 1.0, 2.0, 4.0]
+model:
+  input_size: 16
+  target_size: 32
+  patch_size: 4
+  in_channels: 3
+  embed_dim: 64
+  encoder_depth: 2
+  encoder_heads: 4
+  decoder_dim: 48
+  decoder_depth: 1
+  decoder_heads: 4
+  mask_ratio: 0.75
+  reference_gsd: 1.0
+  blur_kernel: 5
+  band_config:
+    low: {kernel: 5, target: low_frequency}
+    high: {residual: true, target: high_frequency}
+training:
+  epochs: 3
+  batch_size: 8
+  learning_rate: 0.0005
+  weight_decay: 0.05
+  num_workers: 0
+  gradient_accumulation: 1
+  amp: true
+  warmup_fraction: 0.1
+runtime:
+  device: auto
+paths:
+  checkpoint: result/checkpoints/scalemae.pt
+  training_metrics: result/training/metrics.json
+  inference_dir: result/output
+  evaluation_dir: result/evaluation

config.json

@@ -0,0 +1,71 @@
+{
+  "model_name": "Scale-MAE",
+  "model_type": "scale_mae",
+  "architectures": [
+    "ScaleMAE"
+  ],
+  "framework": "PyTorch",
+  "domain": "earth-observation",
+  "task": "multiscale-geospatial-representation-learning",
+  "implementation": {
+    "entry_point": "model/scalemae.py",
+    "scope": "GSD-aware masked autoencoding with paired low- and high-frequency target reconstruction",
+    "train_script": "scripts/train.py",
+    "inference_script": "scripts/inference.py",
+    "evaluation_script": "scripts/result.py",
+    "synthetic_data_script": "scripts/fake_data.py"
+  },
+  "architecture": {
+    "family": "scale-aware masked autoencoder",
+    "input_size": 16,
+    "target_size": 32,
+    "patch_size": 4,
+    "input_patches": 16,
+    "in_channels": 3,
+    "embed_dim": 64,
+    "encoder_depth": 2,
+    "encoder_heads": 4,
+    "decoder_dim": 48,
+    "decoder_depth": 1,
+    "decoder_heads": 4,
+    "mask_ratio": 0.75,
+    "reference_gsd": 1.0,
+    "blur_kernel": 5,
+    "band_config": {
+      "low": {"kernel": 5, "target": "low_frequency"},
+      "high": {"residual": true, "target": "high_frequency"}
+    }
+  },
+  "data": {
+    "datasets": [
+      "FMoW-RGB",
+      "RESISC-45",
+      "EuroSAT",
+      "UCMerced",
+      "AID",
+      "MLRSNet"
+    ],
+    "protocol": "synthetic_multiscale_geospatial_rgb",
+    "format": "NPZ",
+    "input_key": "images",
+    "input_shape": ["N", 3, 16, 16],
+    "target_key": "targets",
+    "target_shape": ["N", 3, 32, 32],
+    "gsd_key": "gsd",
+    "gsd_shape": ["N"],
+    "gsd_values": [0.5, 1.0, 2.0, 4.0],
+    "label_key": "labels",
+    "label_shape": ["N"],
+    "num_classes": 4,
+    "default_train_file": "data/train.npz",
+    "default_test_file": "data/test.npz"
+  },
+  "configuration_sources": [
+    "conf/config.yaml",
+    "model/scalemae.py",
+    "scripts/fake_data.py",
+    "scripts/train.py",
+    "scripts/inference.py",
+    "scripts/result.py"
+  ]
+}

configuration.json

@@ -0,0 +1,12 @@
+{
+  "framework": "PyTorch",
+  "task": "remote_sensing_representation_learning",
+  "model": "Scale-MAE",
+  "input_format": "BCHW_with_GSD",
+  "protocol": "scale_aware_masked_autoencoding",
+  "default_config": "conf/config.yaml",
+  "train": "scripts/train.py",
+  "inference": "scripts/inference.py",
+  "evaluation": "scripts/result.py",
+  "visualization": "scripts/result.py"
+}

model/scalemae.py

@@ -0,0 +1,107 @@
+"""Scale-aware masked autoencoder for paired remote-sensing resolutions."""
+import math
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+def sincos_2d(height, width, dim, device, dtype, scale=1.0):
+    if dim % 4:
+        raise ValueError("embedding dimension must be divisible by four")
+    y, x = torch.meshgrid(torch.arange(height, device=device, dtype=dtype),
+                          torch.arange(width, device=device, dtype=dtype), indexing="ij")
+    quarter = dim // 4
+    omega = torch.exp(-math.log(10000.0) * torch.arange(quarter, device=device, dtype=dtype) /
+                      max(quarter - 1, 1))
+    x, y = x.flatten()[:, None] * scale, y.flatten()[:, None] * scale
+    return torch.cat((torch.sin(x * omega), torch.cos(x * omega),
+                      torch.sin(y * omega), torch.cos(y * omega)), dim=1)
+
+
+class ScaleMAE(nn.Module):
+    def __init__(self, input_size=32, target_size=64, patch_size=4, in_channels=3,
+                 embed_dim=64, encoder_depth=2, encoder_heads=4, decoder_dim=48,
+                 decoder_depth=1, decoder_heads=4, mask_ratio=0.75, reference_gsd=1.0,
+                 blur_kernel=5, band_config=None, **legacy):
+        super().__init__()
+        input_size = legacy.get("image_size", input_size)
+        self.input_size, self.target_size, self.patch_size = input_size, target_size, patch_size
+        self.in_channels, self.embed_dim, self.mask_ratio = in_channels, embed_dim, mask_ratio
+        self.reference_gsd, self.blur_kernel = reference_gsd, blur_kernel
+        if input_size % patch_size or target_size % patch_size:
+            raise ValueError("input_size and target_size must be divisible by patch_size")
+        self.input_grid = input_size // patch_size
+        self.target_grid = target_size // patch_size
+        self.num_patches = self.input_grid ** 2
+        patch_dim = in_channels * patch_size ** 2
+        self.patch_embed = nn.Linear(patch_dim, embed_dim)
+        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
+        self.mask_token = nn.Parameter(torch.zeros(1, 1, decoder_dim))
+        enc = nn.TransformerEncoderLayer(embed_dim, encoder_heads, embed_dim * 4,
+                                         dropout=0, activation="gelu", batch_first=True, norm_first=True)
+        self.encoder, self.encoder_norm = nn.TransformerEncoder(enc, encoder_depth), nn.LayerNorm(embed_dim)
+        self.decoder_input = nn.Linear(embed_dim, decoder_dim)
+        dec = nn.TransformerEncoderLayer(decoder_dim, decoder_heads, decoder_dim * 4,
+                                         dropout=0, activation="gelu", batch_first=True, norm_first=True)
+        self.decoder, self.decoder_norm = nn.TransformerEncoder(dec, decoder_depth), nn.LayerNorm(decoder_dim)
+        self.fpn = nn.Sequential(nn.Conv2d(decoder_dim, decoder_dim, 3, padding=1), nn.GELU(),
+                                 nn.Conv2d(decoder_dim, decoder_dim, 3, padding=1), nn.GELU())
+        self.low_head = nn.Conv2d(decoder_dim, in_channels, 1)
+        self.high_head = nn.Conv2d(decoder_dim, in_channels, 1)
+        self.band_config = band_config or {"low": {"kernel": blur_kernel}, "high": {"residual": True}}
+        nn.init.normal_(self.cls_token, std=.02); nn.init.normal_(self.mask_token, std=.02)
+
+    def patchify(self, images):
+        b, c, h, w = images.shape; p = self.patch_size
+        if (c, h, w) != (self.in_channels, self.input_size, self.input_size):
+            raise ValueError(f"expected input {(self.in_channels, self.input_size, self.input_size)}, got {(c,h,w)}")
+        return images.reshape(b, c, h//p, p, w//p, p).permute(0,2,4,1,3,5).reshape(b, -1, c*p*p)
+
+    def bandpass_targets(self, target):
+        k = int(self.band_config.get("low", {}).get("kernel", self.blur_kernel))
+        if k < 3 or k % 2 == 0: raise ValueError("low-frequency kernel must be odd and at least three")
+        low = F.avg_pool2d(target, k, 1, k//2, count_include_pad=False)
+        return low, target - low
+
+    def _positions(self, gsd, grid, dim):
+        scales = gsd.to(dtype=self.cls_token.dtype).flatten() / self.reference_gsd
+        base = sincos_2d(grid, grid, dim, gsd.device, self.cls_token.dtype)
+        return base.unsqueeze(0) * scales[:, None, None]
+
+    def forward(self, images, gsd, target=None, mask_ratio=None, return_features=True):
+        if target is None: target = F.interpolate(images, (self.target_size, self.target_size), mode="bilinear", align_corners=False)
+        if target.shape[-2:] != (self.target_size, self.target_size):
+            raise ValueError("target resolution does not match target_size")
+        if target.shape[:2] != (images.shape[0], self.in_channels):
+            raise ValueError("target batch or channel dimensions do not match images")
+        b = images.shape[0]; patches = self.patchify(images); n = patches.shape[1]
+        ratio = self.mask_ratio if mask_ratio is None else float(mask_ratio)
+        if not 0 <= ratio < 1: raise ValueError("mask_ratio must be in [0, 1)")
+        noise = torch.rand(b, n, device=images.device)
+        ids_shuffle = noise.argsort(dim=1); ids_restore = ids_shuffle.argsort(dim=1)
+        keep = max(1, int(n * (1 - ratio))); ids_keep = ids_shuffle[:, :keep]
+        visible = torch.gather(self.patch_embed(patches), 1, ids_keep[..., None].expand(-1, -1, self.embed_dim))
+        visible = visible + torch.gather(self._positions(gsd, self.input_grid, self.embed_dim), 1,
+                                         ids_keep[..., None].expand(-1, -1, self.embed_dim))
+        encoded = self.encoder_norm(self.encoder(torch.cat((self.cls_token.expand(b,-1,-1), visible), 1)))
+        cls, visible_encoded = encoded[:, :1], encoded[:, 1:]
+        decoded = self.decoder_input(visible_encoded)
+        full = self.mask_token.to(decoded.dtype).expand(b, n, -1).clone(); full.scatter_(1, ids_keep[..., None].expand(-1,-1,decoded.shape[-1]), decoded)
+        full = full + self._positions(gsd, self.input_grid, full.shape[-1])
+        full = self.decoder_norm(self.decoder(full))
+        fmap = full.transpose(1, 2).reshape(b, -1, self.input_grid, self.input_grid)
+        fmap = F.interpolate(fmap, (self.target_grid, self.target_grid), mode="bilinear", align_corners=False)
+        fmap = self.fpn(fmap)
+        low = F.interpolate(self.low_head(fmap), (self.target_size, self.target_size), mode="bilinear", align_corners=False)
+        high = F.interpolate(self.high_head(fmap), (self.target_size, self.target_size), mode="bilinear", align_corners=False)
+        low_target, high_target = self.bandpass_targets(target)
+        mask = torch.ones(b, n, dtype=torch.bool, device=images.device); mask.scatter_(1, ids_keep, False)
+        pixel_mask = F.interpolate(mask.view(b, 1, self.input_grid, self.input_grid).float(),
+                                   (self.target_size, self.target_size), mode="nearest")
+        masked = pixel_mask.sum().clamp_min(1.0) * self.in_channels
+        low_loss = ((low-low_target).square() * pixel_mask).sum() / masked
+        high_loss = ((high-high_target).square() * pixel_mask).sum() / masked
+        return {"loss": low_loss + high_loss, "low_loss": low_loss, "high_loss": high_loss,
+                "reconstruction": low + high, "low_reconstruction": low, "high_reconstruction": high,
+                "low_target": low_target, "high_target": high_target, "mask": mask,
+                "ids_restore": ids_restore, "features": cls[:, 0] if return_features else None}

scripts/fake_data.py

@@ -0,0 +1,26 @@
+"""Create paired low/high resolution scenes with labels for kNN evaluation."""
+import argparse, json
+from pathlib import Path
+import numpy as np, yaml
+from PIL import Image
+ROOT = Path(__file__).resolve().parents[1]
+
+def make_split(count, cfg, rng):
+    size, target, channels, classes = cfg["input_size"], cfg["target_size"], cfg["channels"], cfg["num_classes"]
+    labels = np.arange(count, dtype=np.int64) % classes; rng.shuffle(labels)
+    gsd = rng.choice(np.asarray(cfg["gsd_values"], dtype=np.float32), count)
+    y, x = np.mgrid[:target, :target].astype(np.float32); images = np.empty((count, channels, size, size), np.float32); targets = np.empty((count, channels, target, target), np.float32)
+    for i, label in enumerate(labels):
+        pattern = np.sin((x + label*2)*np.pi*(label+1)/target) + np.cos((y-label*2)*np.pi*(label+1)/target)
+        pattern = (pattern-pattern.min())/(pattern.max()-pattern.min())
+        scene = np.stack([np.roll(pattern, label*c, axis=c%2) for c in range(channels)])
+        targets[i] = np.clip(scene + rng.normal(0, .02 + .01*gsd[i], scene.shape), 0, 1)
+        images[i] = np.asarray([Image.fromarray((targets[i,c]*255).astype('uint8')).resize((size,size), Image.Resampling.BOX) for c in range(channels)], dtype=np.float32)/255
+    return images, targets, labels, gsd
+
+def main():
+    p=argparse.ArgumentParser(); p.add_argument("--config", default=str(ROOT/"conf/config.yaml")); a=p.parse_args(); cfg=yaml.safe_load(Path(a.config).read_text()); d=cfg["data"]; out=ROOT/d["root"]; out.mkdir(exist_ok=True); rng=np.random.default_rng(cfg["seed"])
+    for split,n in (("train",d["train_samples"]),("test",d["test_samples"])):
+        images,targets,labels,gsd=make_split(n,d,rng); np.savez_compressed(out/f"{split}.npz", images=images, targets=targets, labels=labels, gsd=gsd)
+    (out/"format.json").write_text(json.dumps({"images":"float32 BCHW input resolution","targets":"float32 BCHW target resolution","labels":"int64","gsd":"metres per pixel"},indent=2)+"\n"); print("created",out/"train.npz",out/"test.npz")
+if __name__ == "__main__": main()

scripts/inference.py

@@ -0,0 +1,17 @@
+"""Deterministic feature extraction and paired-resolution reconstruction."""
+import argparse, importlib.util
+from pathlib import Path
+import numpy as np, torch, yaml
+ROOT=Path(__file__).resolve().parents[1]
+def load(cfg):
+    s=importlib.util.spec_from_file_location("scalemae",ROOT/"model/scalemae.py"); m=importlib.util.module_from_spec(s); s.loader.exec_module(m); return m.ScaleMAE(**cfg["model"])
+def run(model,a,device,ratio):
+    with torch.no_grad(): return model(torch.from_numpy(a["images"]).to(device),torch.from_numpy(a["gsd"]).to(device),target=torch.from_numpy(a["targets"]).to(device),mask_ratio=ratio)
+def main():
+    p=argparse.ArgumentParser(); p.add_argument("--config",type=Path,default=ROOT/"conf/config.yaml"); p.add_argument("--data",type=Path); p.add_argument("--train-data",type=Path); p.add_argument("--checkpoint",type=Path); p.add_argument("--output-dir",type=Path); p.add_argument("--device",choices=("auto","cpu","cuda")); a=p.parse_args(); cfg=yaml.safe_load(a.config.read_text()); ck=a.checkpoint or ROOT/cfg["paths"]["checkpoint"]
+    if not ck.exists(): raise FileNotFoundError("Run training before inference")
+    requested=a.device or cfg["runtime"]["device"]
+    if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable")
+    device=torch.device("cuda" if torch.cuda.is_available() and requested!="cpu" else "cpu"); torch.manual_seed(cfg["seed"]); model=load(cfg); model.load_state_dict(torch.load(ck,map_location=device,weights_only=False)["model"]); model.to(device).eval(); train=np.load(a.train_data or ROOT/cfg["data"]["root"]/"train.npz"); test=np.load(a.data or ROOT/cfg["data"]["root"]/"test.npz"); tr=run(model,train,device,0.0); te=run(model,test,device,cfg["model"]["mask_ratio"])
+    out=a.output_dir or ROOT/cfg["paths"]["inference_dir"]; out.mkdir(parents=True,exist_ok=True); np.savez_compressed(out/"reconstruction.npz",target=test["targets"],prediction=te["reconstruction"].cpu().numpy(),low_target=te["low_target"].cpu().numpy(),low_prediction=te["low_reconstruction"].cpu().numpy(),high_target=te["high_target"].cpu().numpy(),high_prediction=te["high_reconstruction"].cpu().numpy(),mask=te["mask"].cpu().numpy(),gsd=test["gsd"],labels=test["labels"],test_features=te["features"].cpu().numpy(),train_features=tr["features"].cpu().numpy(),train_labels=train["labels"],train_gsd=train["gsd"]); print("inference=",out/"reconstruction.npz")
+if __name__=="__main__": main()

scripts/result.py

@@ -0,0 +1,84 @@
+"""Evaluate Scale-MAE reconstruction, kNN transfer and GSD sensitivity."""
+
+import json
+import argparse
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import yaml
+
+ROOT = Path(__file__).resolve().parents[1]
+
+
+def knn_predict(train_features, train_labels, test_features, k=3):
+    train = train_features / np.maximum(np.linalg.norm(train_features, axis=1, keepdims=True), 1e-8)
+    test = test_features / np.maximum(np.linalg.norm(test_features, axis=1, keepdims=True), 1e-8)
+    nearest = np.argsort(-(test @ train.T), axis=1)[:, :min(k, len(train))]
+    return np.asarray([np.bincount(train_labels[row]).argmax() for row in nearest])
+
+
+def rgb(image):
+    return np.clip(image[:3].transpose(1, 2, 0), 0, 1)
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Evaluate Scale-MAE outputs")
+    parser.add_argument("--config", default=str(ROOT / "conf/config.yaml"))
+    args = parser.parse_args()
+    cfg = yaml.safe_load(Path(args.config).read_text())
+    source = ROOT / cfg["paths"]["inference_dir"] / "reconstruction.npz"
+    if not source.exists():
+        raise FileNotFoundError("Run inference before evaluation")
+    data = np.load(source)
+    for key in ("prediction", "target", "test_features", "train_features", "labels", "gsd"):
+        if key not in data: raise ValueError(f"inference archive missing {key}")
+    if not np.isfinite(data["prediction"]).all(): raise FloatingPointError("non-finite inference output")
+    prediction = knn_predict(data["train_features"], data["train_labels"], data["test_features"])
+    accuracy = float(np.mean(prediction == data["labels"]))
+    mse = np.mean((data["prediction"] - data["target"]) ** 2, axis=(1, 2, 3))
+    low_mse = float(np.mean((data["low_prediction"] - data["low_target"]) ** 2))
+    high_mse = float(np.mean((data["high_prediction"] - data["high_target"]) ** 2))
+    gsd_values = sorted(np.unique(data["gsd"]).tolist())
+    gsd_mse = {str(value): float(np.mean(mse[data["gsd"] == value])) for value in gsd_values}
+    gsd_accuracy = {str(value): float(np.mean(prediction[data["gsd"] == value] == data["labels"][data["gsd"] == value]))
+                    for value in gsd_values}
+    result = {"reconstruction_mse": float(np.mean(mse)), "low_frequency_mse": low_mse,
+              "high_frequency_mse": high_mse, "knn_accuracy": accuracy,
+              "gsd_reconstruction_mse": gsd_mse, "gsd_knn_accuracy": gsd_accuracy,
+              "data_source": "synthetic", "protocol": cfg["data"]["protocol"]}
+    output = ROOT / cfg["paths"]["evaluation_dir"]
+    output.mkdir(parents=True, exist_ok=True)
+    (output / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
+
+    figure, axes = plt.subplots(2, 3, figsize=(9, 6))
+    images = ((data["target"][0], "Original"), (data["low_target"][0], "Low target"),
+              (data["high_target"][0] + 0.5, "High target"), (data["prediction"][0], "Reconstruction"),
+              (data["low_prediction"][0], "Low prediction"), (data["high_prediction"][0] + 0.5, "High prediction"))
+    for axis, (image, title) in zip(axes.flat, images):
+        axis.imshow(rgb(image)); axis.set_title(title); axis.axis("off")
+    figure.tight_layout(); figure.savefig(output / "bandpass_reconstruction.png", dpi=160); plt.close(figure)
+
+    figure, axis = plt.subplots(figsize=(6, 3.5))
+    axis.plot(gsd_values, [gsd_mse[str(x)] for x in gsd_values], marker="o", color="#1f77b4")
+    axis.set(xlabel="GSD (m/pixel)", ylabel="Reconstruction MSE", title="Scale-aware Reconstruction")
+    axis.grid(alpha=0.25)
+    figure.tight_layout(); figure.savefig(output / "gsd_reconstruction_error.png", dpi=160); plt.close(figure)
+
+    figure, axis = plt.subplots(figsize=(6, 3.5))
+    axis.plot(gsd_values, [gsd_accuracy[str(x)] for x in gsd_values], marker="s", color="#d62728")
+    axis.set(xlabel="GSD (m/pixel)", ylabel="kNN accuracy", title="Scale-aware Feature Transfer")
+    axis.set_ylim(0, 1.05); axis.grid(alpha=0.25)
+    figure.tight_layout(); figure.savefig(output / "gsd_knn_accuracy.png", dpi=160); plt.close(figure)
+
+    figure, axis = plt.subplots(figsize=(5, 3))
+    axis.bar(["Low frequency", "High frequency"], [low_mse, high_mse], color=["#2a9d8f", "#e76f51"])
+    axis.set(ylabel="MSE", title="Bandpass Reconstruction Error")
+    figure.tight_layout(); figure.savefig(output / "frequency_error.png", dpi=160); plt.close(figure)
+
+    np.save(output / "features.npy", data["test_features"])
+    print(json.dumps(result, indent=2)); print("evaluation=", output)
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,47 @@
+"""Train Scale-MAE with DDP, AMP, accumulation, warmup and cosine decay."""
+import argparse, importlib.util, json, os, random
+from pathlib import Path
+import numpy as np, torch, 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]
+class NPZDataset(Dataset):
+    def __init__(self,path):
+        a=np.load(path); self.images,self.targets,self.gsd=a["images"],a["targets"],a["gsd"]
+    def __len__(self): return len(self.images)
+    def __getitem__(self,i): return torch.from_numpy(self.images[i]), torch.from_numpy(self.targets[i]), torch.tensor(self.gsd[i], dtype=torch.float32)
+def model_class():
+    spec=importlib.util.spec_from_file_location("scalemae",ROOT/"model/scalemae.py"); mod=importlib.util.module_from_spec(spec); spec.loader.exec_module(mod); return mod.ScaleMAE
+def main():
+    p=argparse.ArgumentParser(); p.add_argument("--config",type=Path,default=ROOT/"conf/config.yaml"); p.add_argument("--data",type=Path); p.add_argument("--output",type=Path); p.add_argument("--epochs",type=int); p.add_argument("--batch-size",type=int); p.add_argument("--device",choices=("auto","cpu","cuda")); a=p.parse_args(); cfg=yaml.safe_load(a.config.read_text())
+    if a.epochs is not None: cfg["training"]["epochs"] = a.epochs
+    if a.batch_size is not None: cfg["training"]["batch_size"] = a.batch_size
+    world=int(os.environ.get("WORLD_SIZE","1")); rank=int(os.environ.get("RANK","0")); local=int(os.environ.get("LOCAL_RANK","0")); requested=a.device or cfg["runtime"]["device"]
+    if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA requested but unavailable")
+    cuda=torch.cuda.is_available() and requested!="cpu"; device=torch.device(f"cuda:{local}" if cuda else "cpu")
+    if cuda: torch.cuda.set_device(local)
+    if world>1: dist.init_process_group("nccl" if cuda else "gloo")
+    seed=cfg["seed"]+rank; random.seed(seed); np.random.seed(seed); torch.manual_seed(seed)
+    source=a.data or ROOT/cfg["data"]["root"]/"train.npz"
+    if not source.exists(): raise FileNotFoundError("Run scripts/fake_data.py first")
+    ds=NPZDataset(source); sampler=DistributedSampler(ds,num_replicas=world,rank=rank) if world>1 else None; loader=DataLoader(ds,batch_size=cfg["training"]["batch_size"],shuffle=sampler is None,sampler=sampler,num_workers=cfg["training"]["num_workers"])
+    model=model_class()(**cfg["model"]).to(device); model=DistributedDataParallel(model,device_ids=[local]) if world>1 else model; base=model.module if hasattr(model,"module") else model
+    opt=torch.optim.AdamW(model.parameters(),lr=cfg["training"]["learning_rate"],betas=(.9,.95),weight_decay=cfg["training"]["weight_decay"]); steps=max(1,cfg["training"]["epochs"]*((len(ds)+cfg["training"]["batch_size"]-1)//cfg["training"]["batch_size"])); warm=max(1,int(steps*cfg["training"]["warmup_fraction"])); accum=cfg["training"]["gradient_accumulation"]
+    amp=torch.amp.GradScaler("cuda", enabled=cuda and cfg["training"]["amp"]); history=[]; opt.zero_grad(set_to_none=True); step=0
+    for epoch in range(cfg["training"]["epochs"]):
+        if sampler: sampler.set_epoch(epoch)
+        model.train(); totals=torch.zeros(3,device=device)
+        for batch_idx,(images,targets,gsd) in enumerate(loader):
+            with torch.autocast(device_type="cuda",enabled=cuda and cfg["training"]["amp"]): out=model(images.to(device),gsd.to(device),target=targets.to(device)); loss=out["loss"]/accum
+            if not torch.isfinite(loss): raise FloatingPointError("non-finite training loss")
+            amp.scale(loss).backward(); totals += torch.stack([out["loss"].detach(),out["low_loss"].detach(),out["high_loss"].detach()])
+            if (batch_idx+1)%accum==0 or batch_idx+1==len(loader):
+                amp.unscale_(opt); torch.nn.utils.clip_grad_norm_(model.parameters(),1.0); amp.step(opt); amp.update(); opt.zero_grad(set_to_none=True); step+=1; lr=cfg["training"]["learning_rate"]*(step/warm if step<=warm else .5*(1+np.cos(np.pi*(step-warm)/max(1,steps-warm)))); [g.update(lr=lr) for g in opt.param_groups]
+        if world>1: dist.all_reduce(totals); totals/=world
+        totals/=max(len(loader),1); history.append({"epoch":epoch+1,"loss":totals[0].item(),"low_frequency_loss":totals[1].item(),"high_frequency_loss":totals[2].item()})
+        if rank==0: print(history[-1])
+    if rank==0:
+        ck=a.output or ROOT/cfg["paths"]["checkpoint"]; met=ROOT/cfg["paths"]["training_metrics"]; ck.parent.mkdir(parents=True,exist_ok=True); met.parent.mkdir(parents=True,exist_ok=True); torch.save({"model":base.state_dict(),"config":cfg,"optimizer":opt.state_dict(),"epoch":cfg["training"]["epochs"],"seed":cfg["seed"]},ck); met.write_text(json.dumps({"history":history,"protocol":cfg["data"]["protocol"],"data_source":"synthetic"},indent=2)+"\n"); print("checkpoint=",ck)
+    if world>1: dist.destroy_process_group()
+if __name__=="__main__": main()