""" T-SRDA pretraining — OFFICIAL STCLN protocol. Dataset : 496 fold-5 patches, unlabelled (pretraining_STCLN.py:51 folds=[5]) Batching: 4 PATCHES per batch -> 124 batches/epoch Cropping: inner `for i in range(4): for j in range(4)` loop, each (i, j) a SEPARATE optimizer step -> 124 x 16 = 1,984 crop-steps/epoch Task : NDVI-aware spatiotemporal masking + MSE reconstruction (no labels) """ import os import time import argparse from pathlib import Path import torch import torch.distributed as dist from torch.utils.data import DataLoader, DistributedSampler from torch.cuda.amp import autocast, GradScaler from torch.nn.parallel import DistributedDataParallel as DDP import config as C from dataset import PASTISPatchDataset, pretrain_patch_ids, pad_collate, crop_ij from model import UTAE_ASPP, UTAEPrediction def setup_distributed(): if "RANK" in os.environ: dist.init_process_group(backend="nccl") local_rank = int(os.environ["LOCAL_RANK"]) rank = int(os.environ["RANK"]) world = int(os.environ["WORLD_SIZE"]) torch.cuda.set_device(local_rank) return rank, world, local_rank, torch.device(f"cuda:{local_rank}") return 0, 1, 0, torch.device("cuda:0") def cleanup_distributed(): if dist.is_initialized(): dist.destroy_process_group() def main(): parser = argparse.ArgumentParser() parser.add_argument("--seed", type=int, default=C.SEED, help=f"torch seed; official {C.SEED}, " f"variance runs {C.VARIANCE_SEEDS}") parser.add_argument("--resume", type=str, default="none", help="'auto' = continue from checkpoints/pretrain/" "latest.tar if it exists; 'none' = always start " "fresh; or an explicit .tar path.") args = parser.parse_args() rank, world, local_rank, device = setup_distributed() torch.manual_seed(args.seed) torch.cuda.manual_seed_all(args.seed) # ── Fold-5 patch pool (496 patches) ────────────────────────────────── patch_ids = pretrain_patch_ids() dataset = PASTISPatchDataset(patch_ids, load_target=False) sampler = (DistributedSampler(dataset, num_replicas=world, rank=rank, shuffle=True) if world > 1 else None) train_dl = DataLoader( dataset, batch_size = C.PRE_BATCH, # PATCHES, not crops sampler = sampler, shuffle = (sampler is None), num_workers = 2, pin_memory = True, persistent_workers = True, collate_fn = pad_collate, ) grid = C.PRETRAIN_CROP_GRID crops_per_batch = grid * grid # ── Model ───────────────────────────────────────────────────────────── encoder = UTAE_ASPP( n_channels = C.N_CHANNELS, d_model = C.D_MODEL, n_heads = C.N_HEADS, T_pad = C.T_PAD, windows = C.TSE_WINDOWS, shifts = C.TSE_SHIFTS, ) network = UTAEPrediction( encoder, n_channels = C.N_CHANNELS, d_model = C.D_MODEL, mask_ratio = C.MASK_RATIO, ).to(device) if world > 1: network = DDP(network, device_ids=[local_rank]) optimizer = torch.optim.AdamW(network.parameters(), lr=C.PRE_LR, betas=(0.9, 0.95), weight_decay=C.PRE_WD) scaler = GradScaler() if rank == 0: n_params = sum(p.numel() for p in network.parameters() if p.requires_grad) print(f"[PRETRAIN] patches={len(dataset)} batches/epoch={len(train_dl)} " f"crop-steps/epoch={len(train_dl) * crops_per_batch} " f"world={world} batch={C.PRE_BATCH} patches x {crops_per_batch} crops", flush=True) print(f"[PRETRAIN] seed={args.seed} lr={C.PRE_LR} wd={C.PRE_WD} " f"deep_sup={C.DEEP_SUP} mask_ratio={C.MASK_RATIO} " f"params={n_params/1e6:.2f}M", flush=True) print(f"[PRETRAIN] checkpoints -> {C.CKPT_DIR_PRE}", flush=True) # ── Resume ──────────────────────────────────────────────────────────── os.makedirs(C.CKPT_DIR_PRE, exist_ok=True) latest_path = C.CKPT_DIR_PRE / "latest.tar" start_epoch = 0 resume_from = None if args.resume == "auto": resume_from = latest_path if latest_path.is_file() else None elif args.resume not in ("none", ""): resume_from = Path(args.resume) if not resume_from.is_file(): raise FileNotFoundError(f"--resume: no such checkpoint {resume_from}") if resume_from is not None: ck = torch.load(resume_from, map_location="cpu", weights_only=False) (network.module if world > 1 else network).load_state_dict( ck["model_state_dict"]) optimizer.load_state_dict(ck["optimizer_state_dict"]) if "scaler_state_dict" in ck: scaler.load_state_dict(ck["scaler_state_dict"]) if "torch_rng_state" in ck: torch.set_rng_state(ck["torch_rng_state"].cpu().to(torch.uint8)) if ck.get("cuda_rng_state") is not None and torch.cuda.is_available(): torch.cuda.set_rng_state_all( [s.cpu().to(torch.uint8) for s in ck["cuda_rng_state"]]) start_epoch = int(ck["epoch"]) + 1 if rank == 0: print(f"[PRETRAIN] RESUMED from {resume_from} " f"(finished epoch {ck['epoch']}) -> starting at epoch " f"{start_epoch}/{C.PRE_EPOCHS}", flush=True) if start_epoch >= C.PRE_EPOCHS: print("[PRETRAIN] already complete; nothing to do.", flush=True) cleanup_distributed() return elif rank == 0: print(f"[PRETRAIN] starting fresh at epoch 0 " f"(--resume={args.resume})", flush=True) # ── Training loop ───────────────────────────────────────────────────── for epoch in range(start_epoch, C.PRE_EPOCHS): if sampler is not None: sampler.set_epoch(epoch) network.train() ep_start = time.time() ep_loss = 0.0 n_steps = 0 for it, ((xp, pos, _days), _) in enumerate(train_dl): xp = xp.to(device, non_blocking=True) # (B, T, 10, 128, 128) pos = pos.to(device, non_blocking=True) # (B, T) # inner 4x4 crop loop — one optimizer step per (i, j) for i in range(grid): for j in range(grid): xc, _ = crop_ij(xp, None, i, j, grid) # (B,T,10,32,32) optimizer.zero_grad(set_to_none=True) with autocast(): recon, recon2, target, mask = network(xc, pos) invmask = (1.0 - mask) denom = invmask.sum() + 1e-8 L_main = ((recon - target) ** 2 * invmask).sum() / denom loss = L_main if C.DEEP_SUP: L_deep = ((recon2 - target) ** 2 * invmask).sum() / denom loss = loss + L_deep scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(network.parameters(), max_norm=C.PRE_CLIP) scaler.step(optimizer) scaler.update() ep_loss += loss.item() n_steps += 1 if rank == 0 and it % 10 == 0: el = time.time() - ep_start print(f" ep {epoch:3d} batch {it:4d}/{len(train_dl)} " f"step {n_steps:5d} loss={loss.item():.4f} t={el:.0f}s", flush=True) if rank == 0: ep_time = (time.time() - ep_start) / 60 print(f"[ep {epoch:3d} DONE] mean_loss={ep_loss/n_steps:.4f} " f"steps={n_steps} elapsed={ep_time:.1f}min", flush=True) # Save checkpoints net_to_save = network.module if world > 1 else network # latest.tar EVERY epoch — full resume state. Written to a temp # file then renamed, so a crash mid-write cannot corrupt it. tmp = C.CKPT_DIR_PRE / "latest.tar.tmp" torch.save({ "epoch": epoch, "model_state_dict": net_to_save.state_dict(), "optimizer_state_dict": optimizer.state_dict(), "scaler_state_dict": scaler.state_dict(), "torch_rng_state": torch.get_rng_state(), "cuda_rng_state": (torch.cuda.get_rng_state_all() if torch.cuda.is_available() else None), "seed": args.seed, }, tmp) os.replace(tmp, C.CKPT_DIR_PRE / "latest.tar") # encoder-only companion, so a finetune can start from any epoch tmp_e = C.CKPT_DIR_PRE / "latest.utae.tar.tmp" torch.save({"epoch": epoch, "model_state_dict": net_to_save.utae.state_dict()}, tmp_e) os.replace(tmp_e, C.CKPT_DIR_PRE / "latest.utae.tar") if (epoch == 0 or (epoch + 1) % C.PRE_SAVE_EVERY == 0 or epoch == C.PRE_EPOCHS - 1): ckpt = C.CKPT_DIR_PRE / f"checkpoint_{epoch}.tar" torch.save({ "epoch": epoch, "model_state_dict": net_to_save.state_dict(), "optimizer_state_dict": optimizer.state_dict(), }, ckpt) ckpt_enc = C.CKPT_DIR_PRE / f"checkpoint_{epoch}.utae.tar" torch.save({ "epoch": epoch, "model_state_dict": net_to_save.utae.state_dict(), }, ckpt_enc) print(f"[saved] milestone checkpoint_{epoch}.tar + .utae.tar", flush=True) else: print(f"[saved] latest.tar (epoch {epoch}) — next milestone at " f"epoch {((epoch // C.PRE_SAVE_EVERY) + 1) * C.PRE_SAVE_EVERY - 1}", flush=True) cleanup_distributed() if __name__ == "__main__": main()