Move PASTIS work into PASTIS/ (3518/3518)

0ad6e78 verified 2 days ago

10.8 kB

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

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