scripts/train.py

"""Pre-train SatMAE with masked reconstruction; supports torchrun."""

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, mode):
        archive = np.load(path)
        self.images = archive["images"]
        self.timestamps = archive["timestamps"] if "timestamps" in archive else None
        if mode == "temporal" and self.timestamps is None:
            raise ValueError("temporal datasets must contain timestamps")

    def __len__(self):
        return len(self.images)

    def __getitem__(self, index):
        images = torch.from_numpy(self.images[index])
        if self.timestamps is None:
            return images, torch.empty(0)
        return images, torch.from_numpy(self.timestamps[index])


def load_model_class():
    spec = importlib.util.spec_from_file_location("satmae", ROOT / "model/satmae.py")
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)
    return module.SatMAE


def model_config(config):
    return {
        key: value for key, value in config["model"].items()
        if key not in {"architecture", "runtime_profile"}
    }


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 cosine_learning_rate(progress, config, peak_lr):
    warmup = config["warmup_epochs"]
    if warmup > 0 and progress < warmup:
        return peak_lr * progress / warmup
    span = max(config["epochs"] - warmup, 1)
    phase = min(max((progress - warmup) / span, 0.0), 1.0)
    return config["min_learning_rate"] + 0.5 * (
        peak_lr - config["min_learning_rate"]
    ) * (1.0 + math.cos(math.pi * phase))


def main():
    args = parse_args()
    config = yaml.safe_load(args.config.read_text())
    train_config = config["training"]
    if args.epochs is not None:
        train_config["epochs"] = args.epochs
    if args.batch_size is not None:
        train_config["batch_size"] = args.batch_size

    world_size = int(os.environ.get("WORLD_SIZE", "1"))
    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
    rank = int(os.environ.get("RANK", "0"))
    distributed = world_size > 1
    requested_device = args.device or config["runtime"]["device"]
    use_cuda = torch.cuda.is_available() and requested_device != "cpu"
    if requested_device == "cuda" and not torch.cuda.is_available():
        raise RuntimeError("CUDA was requested but is unavailable")
    if distributed:
        dist.init_process_group("nccl" if use_cuda else "gloo")
    device = torch.device(f"cuda:{local_rank}" if use_cuda else "cpu")
    if use_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"
    if not data_path.exists():
        raise FileNotFoundError(f"training data not found: {data_path}")
    dataset = NPZDataset(data_path, config["model"]["mode"])
    sampler = DistributedSampler(dataset, shuffle=True) if distributed else None
    loader = DataLoader(
        dataset,
        batch_size=train_config["batch_size"],
        shuffle=sampler is None,
        sampler=sampler,
        num_workers=train_config["num_workers"],
        pin_memory=use_cuda,
        drop_last=False,
    )

    model = load_model_class()(**model_config(config)).to(device)
    model_without_ddp = model
    if distributed:
        model = DistributedDataParallel(
            model, device_ids=[local_rank] if use_cuda else None
        )
        model_without_ddp = model.module

    effective_batch = (
        train_config["batch_size"] * train_config["accum_iter"] * world_size
    )
    peak_lr = train_config["learning_rate"]
    if peak_lr is None:
        peak_lr = train_config["base_learning_rate"] * effective_batch / 256
    decay, no_decay = [], []
    for name, parameter in model_without_ddp.named_parameters():
        if not parameter.requires_grad:
            continue
        (no_decay if parameter.ndim == 1 or name.endswith("bias") else decay).append(parameter)
    optimizer = torch.optim.AdamW(
        [
            {"params": decay, "weight_decay": train_config["weight_decay"]},
            {"params": no_decay, "weight_decay": 0.0},
        ],
        lr=peak_lr,
        betas=(0.9, 0.95),
    )
    amp_enabled = bool(config["runtime"].get("amp", True) and use_cuda)
    scaler = torch.amp.GradScaler("cuda", enabled=amp_enabled)
    start_epoch = 0
    history = []
    resume_path = args.resume
    if resume_path is None and train_config.get("resume"):
        resume_path = ROOT / train_config["resume"]
    if resume_path is not None:
        checkpoint = torch.load(resume_path, map_location="cpu", weights_only=False)
        model_without_ddp.load_state_dict(checkpoint["model"])
        optimizer.load_state_dict(checkpoint["optimizer"])
        if checkpoint.get("scaler") is not None:
            scaler.load_state_dict(checkpoint["scaler"])
        start_epoch = checkpoint["epoch"] + 1
        history = 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, train_config["epochs"]):
        if sampler is not None:
            sampler.set_epoch(epoch)
        model.train()
        total_loss = 0.0
        steps = len(loader)
        for step, (images, timestamps) in enumerate(loader):
            progress = epoch + step / max(steps, 1)
            learning_rate = cosine_learning_rate(progress, train_config, peak_lr)
            for group in optimizer.param_groups:
                group["lr"] = learning_rate
            images = images.to(device, non_blocking=use_cuda)
            timestamps = timestamps.to(device, non_blocking=use_cuda)
            timestamps = timestamps if timestamps.numel() else None
            autocast = partial(torch.amp.autocast, "cuda") if amp_enabled else nullcontext
            with autocast():
                output = model(images, timestamps=timestamps)
                loss = output["loss"] / train_config["accum_iter"]
            if not torch.isfinite(loss):
                raise ValueError(f"non-finite loss at epoch {epoch}, step {step}")
            scaler.scale(loss).backward()
            update = (step + 1) % train_config["accum_iter"] == 0 or step + 1 == steps
            if update:
                scaler.step(optimizer)
                scaler.update()
                optimizer.zero_grad(set_to_none=True)
            total_loss += output["loss"].detach().item()

        epoch_loss = total_loss / max(steps, 1)
        record = {
            "epoch": epoch + 1,
            "reconstruction_loss": epoch_loss,
            "learning_rate": optimizer.param_groups[0]["lr"],
        }
        history.append(record)
        if rank == 0:
            print(
                f"epoch={epoch + 1} reconstruction_loss={epoch_loss:.6f} "
                f"lr={record['learning_rate']:.3e}"
            )
            if (epoch + 1) % train_config["save_every"] == 0 or epoch + 1 == train_config["epochs"]:
                checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
                torch.save(
                    {
                        "model": model_without_ddp.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "scaler": scaler.state_dict() if amp_enabled 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" if "synthetic" in data_path.name or (data_path.parent / "format.json").exists() else "provided",
            "effective_batch_size": effective_batch,
            "peak_learning_rate": peak_lr,
        }, indent=2) + "\n")
        print("checkpoint=", checkpoint_path)
    if distributed:
        dist.destroy_process_group()


if __name__ == "__main__":
    main()