"""Pre-train SatMAE++ with masked and native multiscale targets.""" import argparse import importlib.util import json import math import os import random from contextlib import nullcontext from functools import partial 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] class NPZDataset(Dataset): def __init__(self, path, scales, channels, image_size): archive = np.load(path) self.images = archive["images"] expected = (channels, image_size, image_size) if self.images.ndim != 4 or tuple(self.images.shape[1:]) != expected: raise ValueError(f"images must have shape [N, {channels}, {image_size}, {image_size}]") self.targets = { str(scale): archive[f"images_{scale}x"] for scale in scales if scale != 1 and f"images_{scale}x" in archive } missing = [scale for scale in scales if scale != 1 and str(scale) not in self.targets] if missing: raise ValueError(f"dataset is missing native high-resolution targets: {missing}") for scale, values in self.targets.items(): expected_target = (len(self.images), channels, image_size * int(scale), image_size * int(scale)) if tuple(values.shape) != expected_target: raise ValueError(f"images_{scale}x must have shape {expected_target}") def __len__(self): return len(self.images) def __getitem__(self, index): return torch.from_numpy(self.images[index]), { scale: torch.from_numpy(values[index]) for scale, values in self.targets.items() } def load_model_class(): spec = importlib.util.spec_from_file_location("satmaepp", ROOT / "model/satmaepp.py") module = importlib.util.module_from_spec(spec) spec.loader.exec_module(module) return module.SatMAEPP def model_config(config): return {key: value for key, value in config["model"].items() if key not in {"architecture", "runtime_profile"}} def validate_config(config): data, model, training = config["data"], config["model"], config["training"] if data["image_size"] != model["image_size"]: raise ValueError("data.image_size and model.image_size must match") if data["channels"] != model["in_channels"] or data["scales"] != model["scales"]: raise ValueError("data and model channels/scales must match") if model["mode"] == "rgb" and data["channels"] != 3: raise ValueError("RGB protocol requires three channels") if model["mode"] == "multispectral" and data["channels"] != 10: raise ValueError("Sentinel protocol requires ten channels") if training["epochs"] <= 0 or training["batch_size"] <= 0 or training["accum_iter"] <= 0: raise ValueError("epochs, batch_size, and accum_iter must be positive") def validate_dataset_config(config, archive): validate_config(config) data, model = config["data"], config["model"] if "images" not in archive: raise ValueError("dataset must contain images") expected = (data["channels"], data["image_size"], data["image_size"]) if archive["images"].ndim != 4 or tuple(archive["images"].shape[1:]) != expected: raise ValueError(f"images must have shape [N, {expected[0]}, {expected[1]}, {expected[2]}]") def parse_args(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--config", type=Path, default=ROOT / "conf/config.yaml") parser.add_argument("--data", type=Path, default=None) parser.add_argument("--output", type=Path, default=None) parser.add_argument("--resume", type=Path, default=None) parser.add_argument("--epochs", type=int, default=None) parser.add_argument("--batch-size", type=int, default=None) parser.add_argument("--device", choices=("auto", "cpu", "cuda"), default=None) return parser.parse_args() def learning_rate(progress, config, peak): warmup = config["warmup_epochs"] if warmup and progress < warmup: return peak * progress / warmup phase = min(max((progress - warmup) / max(config["epochs"] - warmup, 1), 0), 1) return config["min_learning_rate"] + 0.5 * (peak - config["min_learning_rate"]) * ( 1 + math.cos(math.pi * phase) ) def main(): args = parse_args() config = yaml.safe_load(args.config.read_text()) validate_config(config) training = config["training"] if args.epochs is not None: training["epochs"] = args.epochs if args.batch_size is not None: training["batch_size"] = args.batch_size world = int(os.environ.get("WORLD_SIZE", "1")) rank = int(os.environ.get("RANK", "0")) local_rank = int(os.environ.get("LOCAL_RANK", "0")) requested = args.device or config["runtime"]["device"] cuda = torch.cuda.is_available() and requested != "cpu" if requested == "cuda" and not torch.cuda.is_available(): raise RuntimeError("CUDA was requested but is unavailable") if world > 1: dist.init_process_group("nccl" if cuda else "gloo") device = torch.device(f"cuda:{local_rank}" if cuda else "cpu") if cuda: torch.cuda.set_device(local_rank) seed = config["seed"] + rank random.seed(seed); np.random.seed(seed); torch.manual_seed(seed) data_path = args.data or ROOT / config["data"]["root"] / "train.npz" validate_dataset_config(config, np.load(data_path)) dataset = NPZDataset(data_path, config["model"]["scales"], config["model"]["in_channels"], config["model"]["image_size"]) sampler = DistributedSampler(dataset) if world > 1 else None loader = DataLoader(dataset, batch_size=training["batch_size"], shuffle=sampler is None, sampler=sampler, num_workers=training["num_workers"], pin_memory=cuda) model = load_model_class()(**model_config(config)).to(device) bare_model = model if world > 1: model = DistributedDataParallel(model, device_ids=[local_rank] if cuda else None) bare_model = model.module effective_batch = training["batch_size"] * training["accum_iter"] * world peak_lr = training["learning_rate"] if peak_lr is None: peak_lr = training["base_learning_rate"] * effective_batch / 256 decay, no_decay = [], [] for name, parameter in bare_model.named_parameters(): if parameter.requires_grad: (no_decay if parameter.ndim == 1 or name.endswith("bias") else decay).append(parameter) optimizer = torch.optim.AdamW([ {"params": decay, "weight_decay": training["weight_decay"]}, {"params": no_decay, "weight_decay": 0.0}, ], lr=peak_lr, betas=(0.9, 0.95)) amp = bool(config["runtime"].get("amp", True) and cuda) scaler = torch.amp.GradScaler("cuda", enabled=amp) start_epoch, history = 0, [] resume = args.resume or (ROOT / training["resume"] if training.get("resume") else None) if resume is not None: checkpoint = torch.load(resume, map_location="cpu", weights_only=False) bare_model.load_state_dict(checkpoint["model"]) optimizer.load_state_dict(checkpoint["optimizer"]) if checkpoint.get("scaler"): scaler.load_state_dict(checkpoint["scaler"]) start_epoch, history = checkpoint["epoch"] + 1, checkpoint.get("history", []) checkpoint_path = args.output or ROOT / config["paths"]["checkpoint"] metrics_path = ROOT / config["paths"]["training_metrics"] optimizer.zero_grad(set_to_none=True) for epoch in range(start_epoch, training["epochs"]): if sampler: sampler.set_epoch(epoch) model.train(); totals = {"loss": 0.0, "reconstruction": 0.0, "multiscale": 0.0} for step, (images, targets) in enumerate(loader): lr = learning_rate(epoch + step / max(len(loader), 1), training, peak_lr) for group in optimizer.param_groups: group["lr"] = lr images = images.to(device, non_blocking=cuda) targets = {scale: value.to(device, non_blocking=cuda) for scale, value in targets.items()} autocast = partial(torch.amp.autocast, "cuda") if amp else nullcontext with autocast(): output = model(images, high_resolution_targets=targets) loss = output["loss"] / training["accum_iter"] if not torch.isfinite(loss): raise ValueError("non-finite training loss") scaler.scale(loss).backward() update = (step + 1) % training["accum_iter"] == 0 or step + 1 == len(loader) if update: scaler.step(optimizer); scaler.update(); optimizer.zero_grad(set_to_none=True) totals["loss"] += output["loss"].detach().item() totals["reconstruction"] += output["reconstruction_loss"].detach().item() totals["multiscale"] += output["multiscale_loss"].detach().item() if world > 1: values = torch.tensor([totals[name] for name in ("loss", "reconstruction", "multiscale")], device=device) dist.all_reduce(values, op=dist.ReduceOp.SUM) totals = dict(zip(("loss", "reconstruction", "multiscale"), (value.item() / world for value in values))) record = {"epoch": epoch + 1, "learning_rate": optimizer.param_groups[0]["lr"]} record.update({name + "_loss": value / max(len(loader), 1) for name, value in totals.items()}) history.append(record) if rank == 0: print(f"epoch={epoch + 1} loss={record['loss_loss']:.6f} " f"reconstruction={record['reconstruction_loss']:.6f} multiscale={record['multiscale_loss']:.6f}") if (epoch + 1) % training["save_every"] == 0 or epoch + 1 == training["epochs"]: checkpoint_path.parent.mkdir(parents=True, exist_ok=True) torch.save({"model": bare_model.state_dict(), "optimizer": optimizer.state_dict(), "scaler": scaler.state_dict() if amp else None, "epoch": epoch, "history": history, "config": config}, checkpoint_path) if rank == 0: metrics_path.parent.mkdir(parents=True, exist_ok=True) metrics_path.write_text(json.dumps({"history": history, "protocol": config["data"]["protocol"], "data_source": "synthetic", "effective_batch_size": effective_batch, "peak_learning_rate": peak_lr}, indent=2) + "\n") print("checkpoint=", checkpoint_path) if world > 1: dist.destroy_process_group() if __name__ == "__main__": main()