""" Unified trainer for the data-scaling study. Usage: python train.py --model unet --share 25 python train.py --model unet --share 50 python train.py --model segformer_b0 --share 25 python train.py --model segformer_b0 --share 50 The 100% runs are not trained here — they are bootstrapped from the existing checkpoints in pv_panel_models/{unet,vit}_model/. See bootstrap_100.py. Each run: - reads the matching subset_{share}.txt for training filenames - validates on the full val set every epoch - logs per-epoch metrics (loss, dice, iou, miou, pixel_acc) to logs/{model}_{share}.json - saves a single checkpoint at the highest val Dice: checkpoints/{model}_{share}_best.pth Hyperparameters mirror the existing pv_panel_models/{unet,vit}_model trainers: Adam, lr=1e-4, ReduceLROnPlateau(mode='max', patience=5, factor=0.5), 50 epochs, batch_size=16, image_size=128, CombinedLoss(0.5·BCE + 0.5·Dice), augmentations: HFlip, VFlip, Rot15. """ import argparse import json import os import time from pathlib import Path import torch import torch.optim as optim from torch.utils.data import DataLoader from tqdm import tqdm from dataset import SubsetSolarPanelDataset from metrics import SegMetrics from models import MODEL_REGISTRY THIS_DIR = Path(__file__).resolve().parent REPO_ROOT = THIS_DIR.parents[1] TRAIN_IMG = REPO_ROOT / "final_data" / "train" / "images" TRAIN_MSK = REPO_ROOT / "final_data" / "train" / "masks" VAL_IMG = REPO_ROOT / "final_data" / "val" / "images" VAL_MSK = REPO_ROOT / "final_data" / "val" / "masks" SUBSETS_DIR = THIS_DIR / "subsets" LOG_DIR = THIS_DIR / "logs" CKPT_DIR = THIS_DIR / "checkpoints" def run_epoch(model, loader, criterion, optimizer, device, train: bool): model.train(mode=train) metrics = SegMetrics() total_loss = 0.0 n_batches = 0 desc = "Train" if train else "Val" ctx = torch.enable_grad() if train else torch.no_grad() with ctx: for images, masks in tqdm(loader, desc=desc, leave=False): images = images.to(device, non_blocking=True) masks = masks.to(device, non_blocking=True) if train: optimizer.zero_grad() outputs = model(images) loss = criterion(outputs, masks) if train: loss.backward() optimizer.step() total_loss += loss.item() n_batches += 1 metrics.update(outputs.detach(), masks) avg_loss = total_loss / max(n_batches, 1) return avg_loss, metrics.compute() def parse_args(): p = argparse.ArgumentParser() p.add_argument("--model", required=True, choices=list(MODEL_REGISTRY.keys())) p.add_argument("--share", required=True, type=int, choices=[25, 50]) p.add_argument("--epochs", type=int, default=50) p.add_argument("--batch-size", type=int, default=16) p.add_argument("--image-size", type=int, default=128) p.add_argument("--lr", type=float, default=1e-4) p.add_argument("--num-workers", type=int, default=4) p.add_argument("--seed", type=int, default=42) return p.parse_args() def main(): args = parse_args() torch.manual_seed(args.seed) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") print(f"[run] model={args.model} share={args.share}% device={device}") LOG_DIR.mkdir(parents=True, exist_ok=True) CKPT_DIR.mkdir(parents=True, exist_ok=True) subset_file = SUBSETS_DIR / f"subset_{args.share}.txt" if not subset_file.is_file(): raise FileNotFoundError( f"{subset_file} not found. Run subsets/make_subsets.py first." ) train_set = SubsetSolarPanelDataset( TRAIN_IMG, TRAIN_MSK, file_list=subset_file, image_size=args.image_size, augment=True, ) val_set = SubsetSolarPanelDataset( VAL_IMG, VAL_MSK, file_list=None, image_size=args.image_size, augment=False, ) print(f"[data] train={len(train_set)} val={len(val_set)}") train_loader = DataLoader( train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, ) val_loader = DataLoader( val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True, ) model_fn = MODEL_REGISTRY[args.model] model, criterion = model_fn() model = model.to(device) n_params = sum(p.numel() for p in model.parameters()) print(f"[model] {args.model} params={n_params:,}") optimizer = optim.Adam(model.parameters(), lr=args.lr) scheduler = optim.lr_scheduler.ReduceLROnPlateau( optimizer, mode="max", patience=5, factor=0.5 ) history = { "model": args.model, "share": args.share, "n_train": len(train_set), "n_val": len(val_set), "epochs": [], } best_dice = -1.0 best_epoch = -1 best_path = CKPT_DIR / f"{args.model}_{args.share}_best.pth" log_path = LOG_DIR / f"{args.model}_{args.share}.json" def _fmt(seconds: float) -> str: seconds = int(round(seconds)) h, rem = divmod(seconds, 3600) m, s = divmod(rem, 60) return f"{h:d}:{m:02d}:{s:02d}" if h else f"{m:d}:{s:02d}" t0 = time.time() history["start_time_iso"] = time.strftime("%Y-%m-%dT%H:%M:%S", time.localtime(t0)) for epoch in range(args.epochs): print(f"\nEpoch {epoch + 1}/{args.epochs}") epoch_t0 = time.time() train_t0 = time.time() train_loss, train_m = run_epoch(model, train_loader, criterion, optimizer, device, train=True) train_seconds = time.time() - train_t0 val_t0 = time.time() val_loss, val_m = run_epoch(model, val_loader, criterion, optimizer, device, train=False) val_seconds = time.time() - val_t0 scheduler.step(val_m["dice"]) epoch_seconds = time.time() - epoch_t0 elapsed = time.time() - t0 avg_per_epoch = elapsed / (epoch + 1) eta = avg_per_epoch * (args.epochs - epoch - 1) epoch_record = { "epoch": epoch + 1, "lr": optimizer.param_groups[0]["lr"], "train_loss": train_loss, "val_loss": val_loss, **{f"train_{k}": v for k, v in train_m.items()}, **{f"val_{k}": v for k, v in val_m.items()}, "epoch_seconds": epoch_seconds, "train_seconds": train_seconds, "val_seconds": val_seconds, } history["epochs"].append(epoch_record) print( f" train loss={train_loss:.4f} dice={train_m['dice']:.4f} " f"iou={train_m['iou']:.4f} miou={train_m['miou']:.4f} " f"pixel_acc={train_m['pixel_acc']:.4f}" ) print( f" val loss={val_loss:.4f} dice={val_m['dice']:.4f} " f"iou={val_m['iou']:.4f} miou={val_m['miou']:.4f} " f"pixel_acc={val_m['pixel_acc']:.4f}" ) print( f" time epoch={_fmt(epoch_seconds)} " f"(train={_fmt(train_seconds)} val={_fmt(val_seconds)}) " f"elapsed={_fmt(elapsed)} ETA={_fmt(eta)}" ) # Save logs every epoch (resilient to crashes / SIGINT). with open(log_path, "w") as f: json.dump(history, f, indent=2) if val_m["dice"] > best_dice: best_dice = val_m["dice"] best_epoch = epoch + 1 torch.save({ "epoch": epoch + 1, "model_state_dict": model.state_dict(), "val_metrics": val_m, "model_name": args.model, "share": args.share, }, best_path) print(f" ↳ new best (dice={best_dice:.4f}) → {best_path.name}") total_seconds = time.time() - t0 history["best_epoch"] = best_epoch history["best_val_dice"] = best_dice history["wall_clock_seconds"] = total_seconds history["end_time_iso"] = time.strftime("%Y-%m-%dT%H:%M:%S") with open(log_path, "w") as f: json.dump(history, f, indent=2) print(f"\n[done] best epoch {best_epoch} (dice={best_dice:.4f})") print(f" wall {_fmt(total_seconds)} ({total_seconds:.1f} s)") print(f" best → {best_path}") print(f" log → {log_path}") if __name__ == "__main__": main()