| """Controlled plateau studies using the inherited fixed32 training recipe.""" |
| import argparse |
| import copy |
| import json |
| import math |
| import os |
| import random |
| import time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| 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 |
| from sklearn.metrics import confusion_matrix, cohen_kappa_score |
|
|
| import config as C |
| from dataset import (PASTISPatchDataset, train_patch_ids, val_patch_ids, |
| pad_collate, crop_ij) |
| from model import UTAE_ASPP, UTAEClassificationDual, compute_boundary_target |
| from top5_models import MaskedUTAE_ASPP, MaskedUTAEClassificationDual, SeasonalUTAEClassificationDual |
| from losses import LovaszSoftmaxLoss |
|
|
| PASTIS_NAMES = C.PASTIS_CLASSES |
|
|
|
|
| |
| |
| |
| 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 compute_all_metrics(preds_np, labels_np, num_classes=20): |
| class_ids = list(range(1, num_classes - 1)) |
|
|
| valid = (labels_np > 0) & (labels_np < num_classes - 1) |
| pv = preds_np[valid] |
| lv = labels_np[valid] |
|
|
| if len(lv) == 0: |
| return dict(miou=0.0, oa=0.0, mf1=0.0, kappa=0.0, |
| per_iou={}, per_f1={}) |
|
|
| |
| cm = confusion_matrix(lv, pv, labels=list(range(num_classes))) |
| diag = np.diag(cm)[1:-1].astype(np.float64) |
| sum0 = cm.sum(0)[1:-1].astype(np.float64) |
| sum1 = cm.sum(1)[1:-1].astype(np.float64) |
|
|
| iou = diag / (sum0 + sum1 - diag + 1e-8) |
| f1 = (2 * diag) / (sum0 + sum1 + 1e-8) |
| miou = float(iou.mean()) |
| mf1 = float(f1.mean()) |
| oa = float((pv == lv).mean()) |
| kappa = float(cohen_kappa_score(lv, pv)) |
|
|
| return dict(miou=miou, oa=oa, mf1=mf1, kappa=kappa, |
| per_iou={c: float(iou[i]) for i, c in enumerate(class_ids)}, |
| per_f1={c: float(f1[i]) for i, c in enumerate(class_ids)}) |
|
|
|
|
| def print_metrics(metrics, tag): |
| print(f" [{tag}] mIoU={metrics['miou']:.4f} OA={metrics['oa']:.4f} " |
| f"mF1={metrics['mf1']:.4f} Kappa={metrics['kappa']:.4f}", flush=True) |
| print(f" {'Cls':>4} {'Name':<28} {'IoU':>8} {'F1':>8}", flush=True) |
| print(f" {'-'*44}", flush=True) |
| for c in range(1, C.N_CLASSES - 1): |
| iv = metrics['per_iou'].get(c, 0.0) |
| fv = metrics['per_f1'].get(c, 0.0) |
| name = PASTIS_NAMES.get(c, f"cls{c}") |
| flag = " **DEAD**" if iv < 0.001 else (" *rare*" if iv < 0.10 else "") |
| print(f" {c:>4} {name:<28} {iv:>8.4f} {fv:>8.4f}{flag}", flush=True) |
|
|
|
|
| |
| |
| |
| |
| |
| |
| def aug_dihedral(xc, yc): |
| """One dihedral-8 transform per crop-batch, image and label together.""" |
| if torch.rand(()) > 0.5: |
| xc, yc = xc.flip(-2), yc.flip(-2) |
| if torch.rand(()) > 0.5: |
| xc, yc = xc.flip(-1), yc.flip(-1) |
| k = int(torch.randint(0, 4, ())) |
| if k: |
| xc, yc = torch.rot90(xc, k, (-2, -1)), torch.rot90(yc, k, (-2, -1)) |
| return xc.contiguous(), yc.contiguous() |
|
|
|
|
| def aug_temporal_subsample(xc, pos, p=0.5, drop_max=0.2, keep_min=24): |
| """Drop up to `drop_max` of the frames, order preserved. Positions keep |
| their original index values -> emulates missing acquisitions of an |
| otherwise identical season. No-op with probability 1-p.""" |
| if torch.rand(()) > p: |
| return xc, pos |
| T = xc.shape[1] |
| keep = int(round(T * (1.0 - random.uniform(0.0, drop_max)))) |
| keep = max(keep_min, min(T, keep)) |
| if keep >= T: |
| return xc, pos |
| idx = torch.randperm(T)[:keep].sort().values.to(xc.device) |
| return (xc.index_select(1, idx).contiguous(), |
| pos.index_select(1, idx).contiguous()) |
|
|
|
|
| def aug_frame_dropout(xc, p=0.25, max_frames=2): |
| """Zero 1-2 whole frames (0 == dataset mean after normalisation) — |
| emulates cloudy/lost acquisitions.""" |
| if torch.rand(()) > p: |
| return xc |
| T = xc.shape[1] |
| n = random.randint(1, max_frames) |
| idx = torch.randperm(T)[:n].to(xc.device) |
| xc = xc.clone() |
| xc[:, idx] = 0.0 |
| return xc |
|
|
|
|
| def aug_spectral(xc, gain_std=0.05, bias_std=0.03, noise_std=0.01): |
| """Per-sample, per-band gain/bias jitter + small gaussian noise (data is |
| z-normalised, so these are fractions of one std).""" |
| B, T, Ch, H, W = xc.shape |
| dev = xc.device |
| gain = 1.0 + gain_std * torch.randn(B, 1, Ch, 1, 1, device=dev) |
| bias = bias_std * torch.randn(B, 1, Ch, 1, 1, device=dev) |
| xc = xc * gain + bias |
| if noise_std > 0: |
| xc = xc + noise_std * torch.randn_like(xc) |
| return xc.contiguous() |
|
|
|
|
| |
| |
| |
| class ModelEMA: |
| def __init__(self, model, decay=0.999): |
| self.decay = decay |
| self.updates = 0 |
| raw = model.module if isinstance(model, DDP) else model |
| self.module = copy.deepcopy(raw).eval() |
| for p in self.module.parameters(): |
| p.requires_grad_(False) |
|
|
| @torch.no_grad() |
| def update(self, model): |
| self.updates += 1 |
| d = min(self.decay, (1.0 + self.updates) / (10.0 + self.updates)) |
| raw = model.module if isinstance(model, DDP) else model |
| msd = raw.state_dict() |
| for k, v in self.module.state_dict().items(): |
| if v.dtype.is_floating_point: |
| v.lerp_(msd[k].detach().float(), 1.0 - d) |
| else: |
| v.copy_(msd[k]) |
|
|
|
|
| |
| |
| |
| def build_param_groups(net_raw, base_lr, encoder_mult, wd): |
| buckets = {"enc_dec": [], "enc_nd": [], "head_dec": [], "head_nd": []} |
| for name, p in net_raw.named_parameters(): |
| if not p.requires_grad: |
| continue |
| is_enc = name.startswith("utae.") |
| no_dec = p.ndim <= 1 |
| buckets[("enc_" if is_enc else "head_") + ("nd" if no_dec else "dec") |
| ].append(p) |
| groups = [ |
| {"params": buckets["enc_dec"], "lr0": base_lr * encoder_mult, |
| "weight_decay": wd}, |
| {"params": buckets["enc_nd"], "lr0": base_lr * encoder_mult, |
| "weight_decay": 0.0}, |
| {"params": buckets["head_dec"], "lr0": base_lr, |
| "weight_decay": wd}, |
| {"params": buckets["head_nd"], "lr0": base_lr, |
| "weight_decay": 0.0}, |
| ] |
| return [g for g in groups if len(g["params"]) > 0] |
|
|
|
|
| def lr_factor(step, total_steps, warmup_steps, min_ratio): |
| """Linear warmup then cosine decay to min_ratio.""" |
| if step < warmup_steps: |
| return (step + 1) / max(1, warmup_steps) |
| p = (step - warmup_steps) / max(1, total_steps - warmup_steps) |
| p = min(1.0, p) |
| cos = 0.5 * (1.0 + math.cos(math.pi * p)) |
| return min_ratio + (1.0 - min_ratio) * cos |
|
|
|
|
| def set_lrs(optimizer, factor): |
| for g in optimizer.param_groups: |
| g["lr"] = g["lr0"] * factor |
|
|
|
|
| |
| |
| |
| |
| |
| @torch.no_grad() |
| def validate(net, val_dl, device, crop_size, limit_batches=0): |
| net.eval() |
| grid = max(1, C.PATCH_SIZE // crop_size) |
| positions = [(0, 0), (1, 1)] if grid >= 2 else [(0, 0)] |
| all_preds, all_labels = [], [] |
| for bi, ((xv, pv_, _dv), yv) in enumerate(val_dl): |
| if limit_batches and bi >= limit_batches: |
| break |
| xv = xv.to(device, non_blocking=True) |
| pv_ = pv_.to(device, non_blocking=True) |
| yv = yv.to(device, non_blocking=True) |
| for (i, j) in positions: |
| xc, yc = crop_ij(xv, yv, i, j, grid) |
| with autocast(): |
| out = net(xc, pv_) |
| pred = out["refined_logits"].float().argmax(dim=1) |
| all_preds.append(pred.cpu().numpy().ravel()) |
| all_labels.append(yc.cpu().numpy().ravel()) |
| return compute_all_metrics(np.concatenate(all_preds), |
| np.concatenate(all_labels), C.N_CLASSES) |
|
|
|
|
| def pack_checkpoint(state_dict, epoch, metrics, seed, which, architecture_variant='plain'): |
| return {"epoch": epoch, |
| "model_state_dict": state_dict, |
| "val_miou": metrics["miou"], |
| "val_oa": metrics["oa"], |
| "val_mf1": metrics["mf1"], |
| "val_kappa": metrics["kappa"], |
| "val_per_iou": metrics["per_iou"], |
| "seed": seed, "which": which, "recipe": "plateau_fixed32", |
| "architecture_variant": architecture_variant} |
|
|
|
|
| |
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--pretrain_pth", type=str, required=True) |
| ap.add_argument("--seed", type=int, default=C.SEED) |
| ap.add_argument("--tag", type=str, default="32X32_plateau") |
| ap.add_argument("--out_dir", type=str, default="", |
| help="default: <EXP_ROOT>/checkpoints/finetune_32X32/<tag>") |
| ap.add_argument("--epochs", type=int, default=C.FT_EPOCHS) |
| ap.add_argument("--batch", type=int, default=C.FT_BATCH) |
| |
| ap.add_argument("--crop_size", type=int, default=64, |
| help="random-crop side; 32 = official crop size") |
| ap.add_argument("--crops_per_patch", type=int, default=2, |
| help="random crops per patch per epoch (2 keeps the " |
| "official 76 steps/epoch)") |
| ap.add_argument("--official_crops", action="store_true", |
| help="control run: fixed 32x32 diagonal crops, no augs " |
| "(= v1 recipe with v2 plumbing)") |
| ap.add_argument("--fixed_crops", action="store_true", |
| help="fixed diagonal crops at --crop_size; keeps augmentation settings") |
| ap.add_argument("--freeze_encoder_epochs", type=int, default=0) |
| ap.add_argument("--no_aug", action="store_true") |
| ap.add_argument("--no_temporal_aug", action="store_true") |
| ap.add_argument("--mask_ce_invalid", action="store_true") |
| ap.add_argument("--boundary_weight", type=float, default=0.5) |
| ap.add_argument("--init_finetuned", type=str, default="", |
| help="initialize all weights from fine-tuned checkpoint; reset optimizer and schedule") |
| ap.add_argument("--base_lr", type=float, default=1e-4) |
| ap.add_argument("--encoder_lr_mult", type=float, default=0.5) |
| ap.add_argument("--min_lr_ratio", type=float, default=0.01) |
| ap.add_argument("--warmup_epochs", type=float, default=3.0) |
| ap.add_argument("--wd", type=float, default=0.02) |
| ap.add_argument("--label_smoothing", type=float, default=0.05) |
| ap.add_argument("--ema_decay", type=float, default=0.999) |
| ap.add_argument("--patience", type=int, default=30, |
| help="early stopping on val mIoU; 0 = run all epochs") |
| |
| ap.add_argument("--print_every", type=int, default=5) |
| ap.add_argument("--limit_batches", type=int, default=0, |
| help=">0: truncate train/val loops (smoke test)") |
| ap.add_argument("--resume", type=str, default="auto", |
| help="'auto' resumes from <out_dir>/latest.tar if present; " |
| "'none' starts fresh; or an explicit .tar path") |
| ap.add_argument('--distill_weight',type=float,default=0.0) |
| ap.add_argument('--research', choices=['plain','mask_corrected','class_queries','shared_pool_mask'], required=True) |
| ap.add_argument('--sam_bn_mode', choices=['legacy','batch_no_update'], default='legacy') |
| ap.add_argument('--sam_rho', type=float, default=0.02) |
| ap.add_argument('--seasonal_lr', type=float, default=1e-4) |
| args = ap.parse_args() |
| if args.research == 'swa' and args.resume != 'none': |
| ap.error('SWA requires --resume none until averaging state supports resume') |
|
|
| if args.official_crops: |
| args.crop_size = 32 |
| args.no_aug = True |
|
|
| if args.fixed_crops and args.crop_size not in (32, 64): |
| ap.error("--fixed_crops requires --crop_size 32 or 64") |
|
|
| out_dir = (Path(args.out_dir) if args.out_dir else |
| C.EXP_ROOT / "checkpoints" / "finetune_32X32" / args.tag) |
|
|
| rank, world, local_rank, device = setup_distributed() |
|
|
| |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
| torch.set_num_threads(min(8, os.cpu_count() or 1)) |
| torch.manual_seed(args.seed) |
| torch.cuda.manual_seed_all(args.seed) |
|
|
| train_ids = train_patch_ids() |
| val_ids = val_patch_ids() |
|
|
| train_ds = PASTISPatchDataset(train_ids) |
| val_ds = PASTISPatchDataset(val_ids) |
|
|
| sampler = (DistributedSampler(train_ds, num_replicas=world, rank=rank, |
| shuffle=True) |
| if world > 1 else None) |
| train_dl = DataLoader( |
| train_ds, batch_size=args.batch, sampler=sampler, |
| shuffle=(sampler is None), num_workers=2, pin_memory=True, |
| persistent_workers=True, collate_fn=pad_collate) |
| val_dl = DataLoader( |
| val_ds, batch_size=args.batch, shuffle=False, |
| num_workers=2, pin_memory=True, collate_fn=pad_collate) |
|
|
| n_crops = len(C.FT_CROP_IJ) if (args.official_crops or args.fixed_crops) else args.crops_per_patch |
| steps_per_epoch = len(train_dl) * n_crops |
| total_steps = steps_per_epoch * args.epochs |
| warmup_steps = int(args.warmup_epochs * steps_per_epoch) |
|
|
| if rank == 0: |
| print(f"[FT-32X32] train={len(train_ids)} patches ({len(set(train_ids))} unique) " |
| f"val={len(val_ids)} patches ({len(set(val_ids))} unique)", flush=True) |
| print(f"[FT-32X32] world={world} batch={args.batch} patches x {n_crops} crops " |
| f"-> {steps_per_epoch} steps/epoch/rank", flush=True) |
| print(f"[FT-32X32] crop_size={args.crop_size} " |
| f"({'fixed diagonal' if (args.official_crops or args.fixed_crops) else 'RANDOM'}) " |
| f"aug={'off' if args.no_aug else 'dihedral+temporal+spectral'}", |
| flush=True) |
| print(f"[FT-32X32] base_lr={args.base_lr} enc_mult={args.encoder_lr_mult} " |
| f"wd={args.wd} ls={args.label_smoothing} ema={args.ema_decay} " |
| f"warmup_ep={args.warmup_epochs} min_lr_ratio={args.min_lr_ratio} " |
| f"patience={args.patience}", flush=True) |
| print(f"[FT-32X32] seed={args.seed} epochs={args.epochs} " |
| f"checkpoints -> {out_dir}", flush=True) |
|
|
| |
| from plateau_models import architecture |
| encoder_type, network_type = architecture(args.research) |
| encoder = encoder_type( |
| 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) |
| if rank == 0: |
| print(f"[FT-32X32] loading encoder: {args.pretrain_pth}", flush=True) |
| ckpt = torch.load(args.pretrain_pth, map_location="cpu", weights_only=False) |
| sd = ckpt["model_state_dict"] |
| if not any(k.startswith("spatial_encoder.") for k in sd): |
| |
| sd = {k[len("utae."):]: v for k, v in sd.items() if k.startswith("utae.")} |
| if rank == 0: |
| print("[FT-32X32] (extracted encoder weights from full pretrain " |
| "checkpoint)", flush=True) |
| encoder.load_state_dict(sd, strict=True) |
| anchor = {n: p.detach().clone().to(device) for n,p in encoder.named_parameters()} |
|
|
| network = network_type( |
| encoder, d_model=C.D_MODEL, num_classes=C.N_CLASSES).to(device) |
| net_ref = network |
| if world > 1: |
| network = DDP(network, device_ids=[local_rank]) |
| net_ref = network.module |
|
|
| |
| |
| torch.manual_seed(args.seed + 1000 * rank) |
| random.seed(args.seed + 1000 * rank) |
|
|
| if args.init_finetuned: |
| initial = torch.load(args.init_finetuned, map_location="cpu", weights_only=False) |
| incompatible = net_ref.load_state_dict( |
| initial["model_state_dict"], strict=args.research != 'seasonal') |
| if args.research == 'seasonal': |
| expected = { |
| 'sta.seasonal_scores.weight', 'sta.seasonal_scores.bias', |
| 'sta.seasonal_projection.weight', 'sta.seasonal_projection.bias', |
| } |
| if set(incompatible.missing_keys) != expected or incompatible.unexpected_keys: |
| raise RuntimeError(f'Unexpected seasonal loading differences: {incompatible}') |
| if rank == 0: |
| print(f"[INIT] fine-tuned weights from {args.init_finetuned}; optimizer/schedule start fresh", flush=True) |
|
|
| |
| class_weights = torch.ones(C.N_CLASSES, device=device) |
| class_weights[0] = 0.0 |
| class_weights[C.N_CLASSES - 1] = 0.0 |
|
|
| ce_loss = nn.CrossEntropyLoss(weight=class_weights, |
| label_smoothing=args.label_smoothing) |
| if args.mask_ce_invalid: |
| def masked_ce(logits, target): |
| valid = (target > 0) & (target < C.N_CLASSES - 1) |
| if not valid.any(): |
| return logits.sum() * 0.0 |
| return torch.nn.functional.cross_entropy( |
| logits.permute(0, 2, 3, 1)[valid], target[valid], |
| weight=class_weights, label_smoothing=args.label_smoothing) |
| ce_loss = masked_ce |
| lovasz_loss = LovaszSoftmaxLoss(ignore_indices=[0, C.N_CLASSES - 1]) |
| bce_loss = nn.BCEWithLogitsLoss() |
|
|
| |
| parameter_groups = build_param_groups( |
| net_ref, args.base_lr, args.encoder_lr_mult, args.wd) |
| if args.research == 'seasonal': |
| |
| |
| seasonal_parameters = [(name, p) for name, p in net_ref.named_parameters() |
| if name.startswith('sta.seasonal_')] |
| seasonal_ids = {id(p) for _, p in seasonal_parameters} |
| for group in parameter_groups: |
| group['params'] = [p for p in group['params'] if id(p) not in seasonal_ids] |
| parameter_groups = [g for g in parameter_groups if g['params']] |
| for no_decay in (False, True): |
| selected = [p for _, p in seasonal_parameters if (p.ndim <= 1) == no_decay] |
| if selected: |
| parameter_groups.append({ |
| 'params': selected, 'lr0': args.seasonal_lr, |
| 'weight_decay': 0.0 if no_decay else args.wd, |
| }) |
| optimizer = torch.optim.AdamW( |
| parameter_groups, |
| lr=args.base_lr, betas=(0.9, 0.95)) |
| scaler = GradScaler() |
| ema = ModelEMA(network, decay=args.ema_decay) |
| from research_methods import pixel_contrast |
| features = {} |
| if args.research == 'contrast': |
| net_ref.sta.register_forward_hook(lambda module, inputs, output: features.update(value=output)) |
| swa = torch.optim.swa_utils.AveragedModel(net_ref) if args.research == 'swa' else None |
| best_swa = -1.0 |
|
|
| distill_teacher = None |
| if args.distill_weight: |
| from evaluate_temporal_research import ReferenceTeacher |
| saved_rng = torch.get_rng_state(); saved_cuda_rng = torch.cuda.get_rng_state_all() |
| distill_teacher = ReferenceTeacher(device) |
| torch.set_rng_state(saved_rng); torch.cuda.set_rng_state_all(saved_cuda_rng) |
|
|
| def training_loss(x, positions, target, boundary): |
| with autocast(dtype=torch.bfloat16 if args.research == 'contrast' else torch.float16): |
| outputs = network(x, positions) |
| sem = outputs['sem_logits'].float() |
| refined = outputs['refined_logits'].float() |
| loss = (ce_loss(sem,target) + .75*lovasz_loss(sem,target) |
| + ce_loss(refined,target) + .75*lovasz_loss(refined,target) |
| + args.boundary_weight*bce_loss(outputs['bnd_logits'].float().squeeze(1),boundary)) |
| if distill_teacher is not None: |
| with torch.no_grad(): |
| target_probs = distill_teacher(x, positions) |
| valid = (target > 0) & (target < C.N_CLASSES - 1) |
| if valid.any(): |
| kd = torch.nn.functional.kl_div(refined.log_softmax(1), target_probs, reduction='none').sum(1) |
| loss = loss + args.distill_weight * kd[valid].mean() |
| if args.research == 'l2sp': |
| |
| loss = loss + .001*sum((p-anchor[n]).square().sum() for n,p in net_ref.utae.named_parameters()) |
| elif args.research == 'contrast': |
| loss = loss + .05*pixel_contrast(features['value'], target) |
| elif args.research == 'consistency' and epoch >= 3: |
| with torch.no_grad(), autocast(): |
| teacher = ema.module(aug_spectral(x, gain_std=.02, bias_std=.01, noise_std=.005),positions)['refined_logits'].float().softmax(1) |
| valid = (target>0)&(target<C.N_CLASSES-1)&(teacher.max(1).values>=.8) |
| if valid.any(): |
| kl = torch.nn.functional.kl_div(refined.log_softmax(1),teacher,reduction='none').sum(1) |
| loss = loss + .1*min(1.,(epoch-2)/5)*kl[valid].mean() |
| return loss |
|
|
| if rank == 0: |
| n_params = sum(p.numel() for p in network.parameters() if p.requires_grad) |
| n_enc = sum(p.numel() for n_, p in net_ref.named_parameters() |
| if n_.startswith("utae.")) |
| print(f"[FT-32X32] params={n_params/1e6:.2f}M (encoder {n_enc/1e6:.2f}M)", |
| flush=True) |
|
|
| |
| os.makedirs(out_dir, exist_ok=True) |
| latest_path = out_dir / "latest.tar" |
| start_epoch, best_raw, best_ema = 0, {"miou": -1.0}, {"miou": -1.0} |
| best_monitor, patience_counter, history = -1.0, 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 file {resume_from}") |
|
|
| if resume_from is not None: |
| ck = torch.load(resume_from, map_location="cpu", weights_only=False) |
| net_ref.load_state_dict(ck["model_state_dict"]) |
| ema.module.load_state_dict(ck["ema_state_dict"]) |
| ema.updates = ck.get("ema_updates", 0) |
| optimizer.load_state_dict(ck["optimizer_state_dict"]) |
| scaler.load_state_dict(ck["scaler_state_dict"]) |
| best_raw = ck["best_raw"]; best_ema = ck["best_ema"] |
| best_monitor = ck.get("best_monitor", -1.0) |
| patience_counter = ck.get("patience_counter", 0) |
| history = ck.get("history", []) |
| start_epoch = int(ck["epoch"]) + 1 |
| if "rng_python" in ck: |
| random.setstate(ck["rng_python"]) |
| np.random.set_state(ck["rng_numpy"]) |
| torch.set_rng_state(ck["rng_torch"]) |
| torch.cuda.set_rng_state_all(ck["rng_cuda"]) |
| if rank == 0: |
| print(f"[FT-32X32] RESUMED from {resume_from} (finished epoch " |
| f"{ck['epoch']}) -> epoch {start_epoch}/{args.epochs}", |
| flush=True) |
| if start_epoch >= args.epochs: |
| print("[FT-32X32] already complete; nothing to do.", flush=True) |
| cleanup_distributed() |
| return |
| elif rank == 0: |
| print(f"[FT-32X32] starting fresh (--resume={args.resume})", flush=True) |
|
|
| if rank == 0: |
| print(f"[HYPOTHESIS] temporal_aug={not args.no_temporal_aug} " |
| f"mask_ce_invalid={args.mask_ce_invalid} boundary_weight={args.boundary_weight} " |
| f"ema_decay={args.ema_decay}", flush=True) |
| if args.init_finetuned and start_epoch == 0: |
| m_init = validate(net_ref, val_dl, device, args.crop_size, args.limit_batches) |
| if rank == 0: |
| with open(out_dir / "initial_metrics.json", "w") as f: |
| json.dump(m_init, f, indent=2) |
| print(f"[INITIAL] mIoU={m_init['miou']:.6f}", flush=True) |
| |
| best_raw = {"miou": m_init["miou"], "epoch": -1} |
| best_ema = {"miou": m_init["miou"], "epoch": -1} |
| best_monitor = m_init["miou"] |
| for which in ("raw", "ema"): |
| torch.save(pack_checkpoint(net_ref.state_dict(), -1, m_init, |
| args.seed, which, args.research), |
| out_dir / f"model_best_{which}.tar") |
|
|
| |
| gstep = start_epoch * steps_per_epoch |
| hi = C.PATCH_SIZE - args.crop_size |
|
|
| for epoch in range(start_epoch, args.epochs): |
| if sampler is not None: |
| sampler.set_epoch(epoch) |
| network.train() |
| frozen = epoch < args.freeze_encoder_epochs |
| for parameter in net_ref.utae.parameters(): |
| parameter.requires_grad_(not frozen) |
| if frozen: |
| net_ref.utae.eval() |
| if rank == 0: |
| print(f"[ENCODER] frozen={frozen}", flush=True) |
| ep_start = time.time() |
| ep_loss, n_steps = 0.0, 0 |
| stop_flag = False |
|
|
| for it, ((xp, pos, _days), yp) in enumerate(train_dl): |
| if args.limit_batches and it >= args.limit_batches: |
| break |
| xp = xp.to(device, non_blocking=True) |
| pos = pos.to(device, non_blocking=True) |
| yp = yp.to(device, non_blocking=True) |
|
|
| if args.official_crops or args.fixed_crops: |
| crop_list = [crop_ij(xp, yp, i, j, C.PATCH_SIZE // args.crop_size) |
| for (i, j) in [tuple(x) for x in C.FT_CROP_IJ]] |
| else: |
| crop_list = [] |
| for _ in range(n_crops): |
| i = int(torch.randint(0, hi + 1, ())) |
| j = int(torch.randint(0, hi + 1, ())) |
| crop_list.append( |
| (xp[:, :, :, i:i + args.crop_size, |
| j:j + args.crop_size], |
| yp[:, i:i + args.crop_size, j:j + args.crop_size])) |
|
|
| for xc, yc in crop_list: |
| pos_c = pos |
| if not args.no_aug: |
| xc, yc = aug_dihedral(xc, yc) |
| if not args.no_temporal_aug: |
| xc, pos_c = aug_temporal_subsample(xc, pos_c) |
| xc = aug_frame_dropout(xc) |
| xc = aug_spectral(xc) |
| bnd_gt = compute_boundary_target(yc) |
|
|
| optimizer.zero_grad(set_to_none=True) |
| loss = training_loss(xc,pos_c,yc,bnd_gt) |
| if not torch.isfinite(loss): |
| raise RuntimeError('Nonfinite research loss') |
| scaler.scale(loss).backward() |
| if args.research == 'sam': |
| parameters = [p for p in net_ref.parameters() if p.grad is not None] |
| norm = torch.linalg.vector_norm(torch.stack([p.grad.float().norm()/scaler.get_scale() for p in parameters])) |
| if not torch.isfinite(norm): |
| |
| |
| scaler.unscale_(optimizer) |
| scaler.step(optimizer) |
| scaler.update() |
| optimizer.zero_grad(set_to_none=True) |
| print('[SAM AMP] skipped overflowing first pass; scale=', scaler.get_scale(), flush=True) |
| continue |
| offsets = [] |
| with torch.no_grad(): |
| for parameter in parameters: |
| offset = parameter.grad / scaler.get_scale() * (args.sam_rho/(norm+1e-12)) |
| parameter.add_(offset) |
| offsets.append(offset) |
| optimizer.zero_grad(set_to_none=True) |
| bn = [m for m in net_ref.modules() if isinstance(m, torch.nn.modules.batchnorm._BatchNorm)] |
| states = [m.training for m in bn] |
| momenta = [m.momentum for m in bn] |
| tracked = [m.num_batches_tracked.detach().clone() if m.num_batches_tracked is not None else None for m in bn] |
| if args.sam_bn_mode == 'legacy': |
| for m in bn: m.eval() |
| else: |
| for m in bn: |
| m.train() |
| m.momentum = 0.0 |
| try: |
| second_loss = training_loss(xc,pos_c,yc,bnd_gt) |
| scaler.scale(second_loss).backward() |
| finally: |
| with torch.no_grad(): |
| for parameter,offset in zip(parameters,offsets): parameter.sub_(offset) |
| if args.sam_bn_mode == 'batch_no_update': |
| for m,count in zip(bn,tracked): |
| if count is not None: m.num_batches_tracked.copy_(count) |
| for m,state,momentum in zip(bn,states,momenta): |
| m.train(state) |
| m.momentum = momentum |
| scaler.unscale_(optimizer) |
| torch.nn.utils.clip_grad_norm_(network.parameters(),max_norm=5.0) |
| factor = 1.0 if args.research == 'swa' else lr_factor(gstep,total_steps,warmup_steps,args.min_lr_ratio) |
| set_lrs(optimizer,factor) |
| scaler.step(optimizer) |
| scaler.update() |
| ema.update(network) |
| gstep += 1 |
|
|
| ep_loss += loss.item() |
| n_steps += 1 |
|
|
| if rank == 0 and it % 10 == 0: |
| print(f" ep {epoch:3d} batch {it:3d}/{len(train_dl)} " |
| f"step {n_steps:3d}/{steps_per_epoch} " |
| f"loss={loss.item():.4f} " |
| f"lr={optimizer.param_groups[-1]['lr']:.2e} " |
| f"t={time.time()-ep_start:.0f}s", flush=True) |
|
|
| if swa is not None and epoch >= 5: |
| swa.update_parameters(net_ref) |
| |
| swa.module.eval() |
| bn = [m for m in swa.module.modules() if isinstance(m, torch.nn.modules.batchnorm._BatchNorm)] |
| momenta = [m.momentum for m in bn] |
| for m in bn: |
| m.reset_running_stats(); m.momentum = None; m.train() |
| with torch.no_grad(): |
| for bi, ((xx,pp,_), yy) in enumerate(train_dl): |
| if args.limit_batches and bi >= args.limit_batches: break |
| xx,pp,yy = xx.to(device),pp.to(device),yy.to(device) |
| for i,j in C.FT_CROP_IJ: |
| crop,_ = crop_ij(xx,yy,i,j,C.PATCH_SIZE//args.crop_size) |
| with autocast(): swa.module(crop,pp) |
| for m,momentum in zip(bn,momenta): m.momentum=momentum; m.eval() |
|
|
| |
| if rank == 0: |
| m_raw = validate(net_ref, val_dl, device, args.crop_size, |
| args.limit_batches) |
| m_ema = validate(ema.module, val_dl, device, args.crop_size, |
| args.limit_batches) |
| m_swa = validate(swa.module,val_dl,device,args.crop_size,args.limit_batches) if swa is not None and epoch >= 5 else None |
| if m_swa is not None and m_swa['miou'] > best_swa: |
| best_swa = m_swa['miou'] |
| torch.save(pack_checkpoint(swa.module.state_dict(),epoch,m_swa,args.seed,'swa'),out_dir/'model_best_swa.tar') |
| ep_sec = time.time() - ep_start |
| mean_loss = ep_loss / max(n_steps, 1) |
|
|
| history.append({ |
| "epoch": epoch, "train_loss": mean_loss, |
| "lr": optimizer.param_groups[-1]["lr"], |
| "raw": {k: m_raw[k] for k in ("miou", "oa", "mf1", "kappa")}, |
| "ema": {k: m_ema[k] for k in ("miou", "oa", "mf1", "kappa")}, |
| "epoch_sec": ep_sec, "swa": m_swa, "research": args.research, |
| }) |
| with open(out_dir / "history.json", "w") as f: |
| json.dump(history, f, indent=1) |
|
|
| improved = False |
| if m_raw["miou"] > best_raw["miou"]: |
| best_raw = {"miou": m_raw["miou"], "epoch": epoch} |
| torch.save(pack_checkpoint(net_ref.state_dict(), epoch, |
| m_raw, args.seed, "raw", |
| args.research), |
| out_dir / "model_best_raw.tar") |
| improved = True |
| if m_ema["miou"] > best_ema["miou"]: |
| best_ema = {"miou": m_ema["miou"], "epoch": epoch} |
| torch.save(pack_checkpoint(ema.module.state_dict(), epoch, |
| m_ema, args.seed, "ema", |
| args.research), |
| out_dir / "model_best_ema.tar") |
| improved = True |
|
|
| monitor = max(m_raw["miou"], m_ema["miou"]) |
| if monitor > best_monitor: |
| best_monitor = monitor |
| patience_counter = 0 |
| else: |
| patience_counter += 1 |
|
|
| print(f"[ep {epoch:3d} DONE] loss={mean_loss:.4f} " |
| f"raw_mIoU={m_raw['miou']:.4f} ema_mIoU={m_ema['miou']:.4f} " |
| f"best_raw={best_raw['miou']:.4f}@{best_raw['epoch']} " |
| f"best_ema={best_ema['miou']:.4f}@{best_ema['epoch']} " |
| f"patience={patience_counter}/{args.patience} " |
| f"t={ep_sec:.0f}s", flush=True) |
| if improved or (epoch % args.print_every == 0): |
| print_metrics(m_raw, f"ep {epoch} RAW") |
| print_metrics(m_ema, f"ep {epoch} EMA") |
|
|
| |
| tmp = out_dir / "latest.tar.tmp" |
| torch.save({ |
| "epoch": epoch, |
| "model_state_dict": net_ref.state_dict(), |
| "ema_state_dict": ema.module.state_dict(), |
| "ema_updates": ema.updates, |
| "optimizer_state_dict": optimizer.state_dict(), |
| "scaler_state_dict": scaler.state_dict(), |
| "best_raw": best_raw, "best_ema": best_ema, |
| "best_monitor": best_monitor, |
| "patience_counter": patience_counter, |
| "history": history, |
| "args": vars(args), |
| "rng_python": random.getstate(), |
| "rng_numpy": np.random.get_state(), |
| "rng_torch": torch.get_rng_state(), |
| "rng_cuda": torch.cuda.get_rng_state_all(), |
| }, tmp) |
| os.replace(tmp, latest_path) |
|
|
| if args.patience > 0 and patience_counter >= args.patience: |
| print(f"[EARLY STOP] best_raw={best_raw} best_ema={best_ema}", |
| flush=True) |
| stop_flag = True |
|
|
| |
| if world > 1: |
| flag = torch.tensor([1 if stop_flag else 0], device=device) |
| dist.broadcast(flag, src=0) |
| stop_flag = bool(flag.item()) |
| if stop_flag: |
| break |
|
|
| if rank == 0: |
| print(f"\n[FT-v2 COMPLETE] best_raw={best_raw} best_ema={best_ema}", |
| flush=True) |
| print(f"[FT-32X32] checkpoints in {out_dir}", flush=True) |
| cleanup_distributed() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
