| """ |
| Train one baseline / ablation variant with the shared protocol (experiments/configs.py). |
| |
| python -m experiments.train --variant s2f_full |
| python -m experiments.train --variant unet --seed 1 |
| python -m experiments.train --variant s2f_no_gan --resume # continue an interrupted run |
| python -m experiments.train --variant unet --epochs 1 --max-batches 5 # smoke test (1 epoch per phase) |
| |
| Two phases (see experiments/configs.py). Outputs in experiments/runs/<variant>_seed<seed>/: |
| config.json exact configuration used |
| phase1/, phase2/ per phase: history.json (every loss term, val loss / WFM-RME / pixel r, lr, time), |
| best_checkpoint.pth (lowest validation loss of that phase), last_checkpoint.pth |
| (full state for --resume), visualizations/ (first validation batch every 10 epochs) |
| best_checkpoint.pth final model = best phase-2 generator (written when the run is complete) |
| history.json both phases concatenated (written when the run is complete) |
| The training loop follows training/s2f_trainer.py (same losses, label convention and update order); |
| it additionally skips the discriminator when lambda_gan = 0, records every loss term, and seeds all RNGs. |
| """ |
| import argparse |
| import json |
| import math |
| import os |
| import random |
| import shutil |
| import time |
|
|
| import numpy as np |
| import torch |
| import torch.nn as nn |
| import torch.nn.functional as F |
|
|
| from experiments.ablation_models import build_discriminator, build_generator, n_params |
| from experiments.configs import get_config |
| from experiments.layout import DEFAULT_DATASET_ROOT, RUNS_DIR, S2F_ROOT |
| from experiments.losses import REC_LOSSES, PerImageWFMRMELoss |
| from experiments.manifest_data import make_loader |
| from models.s2f_model import create_settings_channels |
| from utils.substrate_settings import compute_settings_normalization |
|
|
| CONFIG_PATH = os.path.join(S2F_ROOT, "config", "substrate_settings.json") |
|
|
|
|
| def seed_everything(seed): |
| random.seed(seed); np.random.seed(seed); torch.manual_seed(seed); torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def cosine_with_warmup(optimizer, warmup_steps, total_steps): |
| """Same schedule as diffusers.optimization.get_cosine_schedule_with_warmup (num_cycles=0.5).""" |
| def f(step): |
| if step < warmup_steps: |
| return step / max(1, warmup_steps) |
| progress = (step - warmup_steps) / max(1, total_steps - warmup_steps) |
| return max(0.0, 0.5 * (1.0 + math.cos(math.pi * progress))) |
| return torch.optim.lr_scheduler.LambdaLR(optimizer, f) |
|
|
|
|
| def add_settings(x, batch, cfg, norm, device): |
| if cfg["substrate_mode"] == "none": |
| return x |
| meta = batch[4] |
| return torch.cat([x, create_settings_channels(meta, norm, device, x.shape, config_path=CONFIG_PATH)], dim=1) |
|
|
|
|
| def save_vis(G, val_loader, cfg, norm, device, path): |
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| batch = next(iter(val_loader)) |
| x = add_settings(batch[0].to(device), batch, cfg, norm, device) |
| with torch.no_grad(): |
| p = (G(x) + 1) / 2 |
| n = min(2, x.size(0)) |
| fig, ax = plt.subplots(3, n, figsize=(4 * n, 12), squeeze=False) |
| for i in range(n): |
| ax[0, i].imshow(x[i, 0].cpu(), cmap="gray"); ax[1, i].imshow(p[i, 0].cpu(), cmap="jet", vmin=0, vmax=1) |
| ax[2, i].imshow(batch[1][i, 0], cmap="jet", vmin=0, vmax=1) |
| for r in range(3): |
| ax[r, i].axis("off") |
| fig.tight_layout(); fig.savefig(path, dpi=100); plt.close(fig) |
|
|
|
|
| def pixel_r(p, y): |
| """Mean per-image Pearson correlation between predicted and true maps (monitoring only).""" |
| p, y = p.flatten(1), y.flatten(1) |
| p, y = p - p.mean(1, keepdim=True), y - y.mean(1, keepdim=True) |
| return ((p * y).sum(1) / (p.norm(dim=1) * y.norm(dim=1) + 1e-12)).mean() |
|
|
|
|
| def _done(phase_dir, epochs): |
| h = os.path.join(phase_dir, "history.json") |
| return (os.path.exists(h) and len(json.load(open(h))) >= epochs |
| and os.path.exists(os.path.join(phase_dir, "best_checkpoint.pth"))) |
|
|
|
|
| def run_phase(G, ph, phase_dir, cfg, train_loader, val_loader, norm, device, resume, max_batches): |
| """One training phase. Fresh discriminator, optimisers and schedule (as in the original staged training); |
| the generator continues from its current weights. Checkpoint selection: lowest validation loss of this phase.""" |
| os.makedirs(os.path.join(phase_dir, "visualizations"), exist_ok=True) |
| use_gan = cfg["lambda_gan"] > 0 |
| D = build_discriminator(cfg["in_channels"]).to(device) if use_gan else None |
| g_opt = torch.optim.Adam(G.parameters(), lr=ph["g_lr"], betas=(cfg["beta1"], cfg["beta2"])) |
| d_opt = torch.optim.Adam(D.parameters(), lr=ph["d_lr"], betas=(cfg["beta1"], cfg["beta2"])) if use_gan else None |
| n_batches = min(len(train_loader), max_batches) if max_batches else len(train_loader) |
| steps = n_batches * ph["epochs"] |
| g_sch = cosine_with_warmup(g_opt, int(steps * cfg["warmup_frac"]), steps) |
| d_sch = cosine_with_warmup(d_opt, int(steps * cfg["warmup_frac"]), steps) if use_gan else None |
| rec_loss, rme, bce = REC_LOSSES[ph["rec_loss"]](), PerImageWFMRMELoss(), nn.BCEWithLogitsLoss() |
| hist, best, start = [], float("inf"), 0 |
| last_path = os.path.join(phase_dir, "last_checkpoint.pth") |
| if resume and os.path.exists(last_path): |
| ck = torch.load(last_path, map_location="cpu", weights_only=False) |
| G.load_state_dict(ck["generator_state_dict"]); g_opt.load_state_dict(ck["g_optimizer"]); g_sch.load_state_dict(ck["g_scheduler"]) |
| if use_gan: |
| D.load_state_dict(ck["discriminator_state_dict"]); d_opt.load_state_dict(ck["d_optimizer"]); d_sch.load_state_dict(ck["d_scheduler"]) |
| hist, best, start = ck["history"], ck["best_val_loss"], ck["epoch"] + 1 |
| print(f"[{ph['name']}] resumed at epoch {start}") |
| print(f"[{ph['name']}] {ph['epochs']} epochs, loss {ph['rec_loss']} x{ph['lambda_rec']:g}, " |
| f"GAN x{cfg['lambda_gan']:g}, force x{ph['lambda_force']:g}, lr {ph['g_lr']:g}", flush=True) |
|
|
| for epoch in range(start, ph["epochs"]): |
| t0 = time.time(); G.train(); D.train() if use_gan else None |
| acc = dict(g=0.0, rec=0.0, gan=0.0, force=0.0, d=0.0); nb = 0 |
| for bi, batch in enumerate(train_loader): |
| if max_batches and bi >= max_batches: |
| break |
| x = add_settings(batch[0].to(device), batch, cfg, norm, device) |
| y = batch[1].to(device) |
| fake = G(x); p = (fake + 1) / 2 |
| l_rec = rec_loss(p, y) |
| l_force = F.mse_loss(p.flatten(1).mean(1), y.flatten(1).mean(1)) |
| l_gan = torch.zeros((), device=device) |
| if use_gan: |
| pf = D(torch.cat([x, fake], 1)); l_gan = bce(pf, torch.ones_like(pf)) |
| l_g = ph["lambda_rec"] * l_rec + cfg["lambda_gan"] * l_gan + ph["lambda_force"] * l_force |
| g_opt.zero_grad(); l_g.backward(); g_opt.step() |
| l_d = torch.zeros(()) |
| if use_gan: |
| pr = D(torch.cat([x, y * 2 - 1], 1)); pf = D(torch.cat([x, fake.detach()], 1)) |
| l_d = 0.5 * (bce(pr, torch.ones_like(pr)) + bce(pf, torch.zeros_like(pf))) |
| d_opt.zero_grad(); l_d.backward(); d_opt.step(); d_sch.step() |
| g_sch.step() |
| for k, v in zip(acc, [l_g, l_rec, l_gan, l_force, l_d]): |
| acc[k] += float(v.detach()) |
| nb += 1 |
| G.eval(); v = dict(rec=0.0, rme=0.0, mse=0.0, r=0.0); nv = 0 |
| with torch.no_grad(): |
| for bi, batch in enumerate(val_loader): |
| if max_batches and bi >= max_batches: |
| break |
| x = add_settings(batch[0].to(device), batch, cfg, norm, device); y = batch[1].to(device) |
| p = (G(x) + 1) / 2 |
| v["rec"] += float(rec_loss(p, y)); v["rme"] += float(rme(p, y)) |
| v["mse"] += float(F.mse_loss(p, y)); v["r"] += float(pixel_r(p, y)); nv += 1 |
| h = {"phase": ph["name"], "epoch": epoch, **{f"train_{k}": a / max(nb, 1) for k, a in acc.items()}, |
| **{f"val_{k}": a / max(nv, 1) for k, a in v.items()}, |
| "g_lr": g_opt.param_groups[0]["lr"], "epoch_seconds": time.time() - t0} |
| hist.append(h) |
| is_best = h["val_rec"] < best |
| best = min(best, h["val_rec"]) |
| print(f"[{ph['name']}] epoch {epoch:3d} G {h['train_g']:.4f} (rec {h['train_rec']:.4f} gan {h['train_gan']:.4f} " |
| f"force {h['train_force']:.5f}) D {h['train_d']:.4f} val_rec {h['val_rec']:.4f} " |
| f"val_r {h['val_r']:.3f} val_rme {h['val_rme']:.3f} {'*' if is_best else ' '} {h['epoch_seconds']:.0f}s", flush=True) |
| if is_best: |
| torch.save({"epoch": epoch, "phase": ph["name"], "generator_state_dict": G.state_dict(), "config": cfg, |
| "val_loss": h["val_rec"], "history": hist}, os.path.join(phase_dir, "best_checkpoint.pth")) |
| state = {"epoch": epoch, "generator_state_dict": G.state_dict(), "g_optimizer": g_opt.state_dict(), |
| "g_scheduler": g_sch.state_dict(), "history": hist, "best_val_loss": best, "config": cfg} |
| if use_gan: |
| state.update(discriminator_state_dict=D.state_dict(), d_optimizer=d_opt.state_dict(), d_scheduler=d_sch.state_dict()) |
| torch.save(state, last_path) |
| json.dump(hist, open(os.path.join(phase_dir, "history.json"), "w"), indent=1) |
| if epoch % 10 == 0 or epoch == ph["epochs"] - 1: |
| save_vis(G, val_loader, cfg, norm, device, os.path.join(phase_dir, "visualizations", f"epoch_{epoch:03d}.png")) |
| return hist |
|
|
|
|
| def train(cfg, run_dir, dataset_root=DEFAULT_DATASET_ROOT, device=None, resume=False, max_batches=None, |
| num_workers=2, allow_cpu=False): |
| device = device or ("cuda" if torch.cuda.is_available() else "cpu") |
| if str(device).startswith("cpu") and not allow_cpu: |
| raise SystemExit("CUDA is not available - refusing to train on CPU (pass --allow-cpu for a smoke test)") |
| if str(device).startswith("cuda"): |
| torch.backends.cudnn.benchmark = True |
| print(f"GPU: {torch.cuda.get_device_name(torch.device(device))} " |
| f"(CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', 'all')})") |
| os.makedirs(run_dir, exist_ok=True) |
| seed_everything(cfg["seed"]) |
| train_loader = make_loader(cfg["train_manifest"], dataset_root, cfg["batch_size"], cfg["img_size"], |
| augment=cfg["augment"], shuffle=True, substrate_mode=cfg["substrate_mode"], |
| input_mode=cfg["input_mode"], seed=cfg["seed"], num_workers=num_workers) |
| if num_workers > 0: |
| try: |
| next(iter(train_loader)) |
| except RuntimeError as e: |
| print(f"DataLoader workers failed ({str(e).splitlines()[0][:120]}); falling back to num_workers=0") |
| train_loader = make_loader(cfg["train_manifest"], dataset_root, cfg["batch_size"], cfg["img_size"], |
| augment=cfg["augment"], shuffle=True, substrate_mode=cfg["substrate_mode"], |
| input_mode=cfg["input_mode"], seed=cfg["seed"], num_workers=0) |
| seed_everything(cfg["seed"]) |
| val_loader = make_loader(cfg["val_manifest"], dataset_root, cfg["batch_size"], cfg["img_size"], |
| substrate_mode=cfg["substrate_mode"], input_mode=cfg["input_mode"]) |
| norm = compute_settings_normalization(config_path=CONFIG_PATH) |
| G = build_generator(**cfg).to(device) |
| n_d = n_params(build_discriminator(cfg["in_channels"])) if cfg["lambda_gan"] > 0 else 0 |
| cfg = dict(cfg, n_params_generator=n_params(G), n_params_discriminator=n_d, |
| n_train=len(train_loader.dataset), n_val=len(val_loader.dataset), device=str(device), |
| torch_version=torch.__version__) |
| json.dump(cfg, open(os.path.join(run_dir, "config.json"), "w"), indent=2) |
| print(f"[{cfg['variant']}] G {cfg['n_params_generator']/1e6:.2f} M params, D {n_d/1e6:.2f} M; " |
| f"train {cfg['n_train']}, val {cfg['n_val']}, device {device}", flush=True) |
|
|
| all_hist = [] |
| for ph in cfg["phases"]: |
| phase_dir = os.path.join(run_dir, ph["name"]) |
| shared = cfg.get("share_phase1_with") |
| if ph["name"] == "phase1" and shared and not _done(phase_dir, ph["epochs"]): |
| src = os.path.join(os.path.dirname(run_dir), f"{shared}_seed{cfg['seed']}", "phase1") |
| if _done(src, ph["epochs"]): |
| print(f"[phase1] identical to {shared} (same seed, data and phase-1 settings) - copying {src}") |
| shutil.copytree(src, phase_dir, dirs_exist_ok=True) |
| else: |
| print(f"[phase1] {shared} phase 1 not finished yet - training phase 1 for this run") |
| if _done(phase_dir, ph["epochs"]) and (resume or shared): |
| print(f"[{ph['name']}] already finished - skipping") |
| hist = json.load(open(os.path.join(phase_dir, "history.json"))) |
| else: |
| hist = run_phase(G, ph, phase_dir, cfg, train_loader, val_loader, norm, device, resume, max_batches) |
| all_hist += hist |
| best_ck = torch.load(os.path.join(phase_dir, "best_checkpoint.pth"), map_location="cpu", weights_only=False) |
| G.load_state_dict(best_ck["generator_state_dict"]) |
| best_ck["config"] = cfg |
| torch.save(best_ck, os.path.join(run_dir, "best_checkpoint.pth")) |
| json.dump(all_hist, open(os.path.join(run_dir, "history.json"), "w"), indent=1) |
| print(f"done. final model: {best_ck['phase']} epoch {best_ck['epoch']}, val loss {best_ck['val_loss']:.4f}; in {run_dir}") |
| return run_dir |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) |
| ap.add_argument("--variant", required=True) |
| ap.add_argument("--seed", type=int, default=None) |
| ap.add_argument("--epochs", type=int, default=None) |
| ap.add_argument("--g-lr", type=float, default=None); ap.add_argument("--d-lr", type=float, default=None) |
| ap.add_argument("--batch-size", type=int, default=None) |
| ap.add_argument("--dataset-root", default=DEFAULT_DATASET_ROOT) |
| ap.add_argument("--run-dir", default=None) |
| ap.add_argument("--device", default=None) |
| ap.add_argument("--resume", action="store_true") |
| ap.add_argument("--max-batches", type=int, default=None, help="limit batches per epoch (smoke test only)") |
| ap.add_argument("--num-workers", type=int, default=2, help="DataLoader worker processes for augmentation") |
| ap.add_argument("--allow-cpu", action="store_true", help="allow training without CUDA (smoke tests)") |
| a = ap.parse_args() |
| cfg = get_config(a.variant, seed=a.seed, epochs=a.epochs, g_lr=a.g_lr, d_lr=a.d_lr, batch_size=a.batch_size) |
| run_dir = a.run_dir or os.path.join(RUNS_DIR, f"{a.variant}_seed{cfg['seed']}") |
| train(cfg, run_dir, a.dataset_root, a.device, a.resume, a.max_batches, a.num_workers, a.allow_cpu) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
