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