scripts/train.py

"""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()