experiments: two-phase training protocol

.gitignore

@@ -13,4 +13,4 @@ S2FApp/ckp/spheroid/ckp_spheroid_GN.pth
 S2FApp/*.docx
 
 # experiments: trained ablation checkpoints (large)
-experiments/runs/
+experiments/runs*/

experiments/README.md

@@ -13,7 +13,7 @@ experiments/
   prepare_dataset.py  audit legacy data vs dataset/, copy missing images, write splits/*.csv
   manifest_data.py    data loader driven by a split manifest (same preprocessing as data/cell_dataset.py)
   ablation_models.py  plain U-Net and S2F generator variants
-  losses.py           per-image WFM-RME reconstruction loss (as used for the released model)
+  losses.py           per-image WFM-RME and weighted MSE + soft Dice losses (both used for the released model)
   configs.py          shared training protocol + one entry per variant
   train.py            train one variant
   evaluate.py         evaluate released checkpoints / baselines / trained runs on all manifests
@@ -54,8 +54,8 @@ python -m experiments.run_all --legacy-root <path>/cellforce/dataset --steps pre
 # 3. released checkpoints, substrate sweep, input-free baseline, image-only stage-1 model
 python -m experiments.run_all --steps eval_released --stage1 <path>/outputs/checkpoints/last/last_checkpoint_final.pth
 
-# 4. baselines / ablations (about 7-8 h each on an 11 GB GPU; resumable), then their evaluation
-python -m experiments.run_all --steps train,eval_runs --variants s2f_full,unet,s2f_no_gan,mask_only,s2f_no_substrate
+# 4. baselines / ablations (S2F about 7 h, U-Net about 2.5 h each on an RTX A6000; resumable), then their evaluation
+python -m experiments.run_all --steps train,eval_runs --variants s2f_full,unet,unet_gan,s2f_no_gan,s2f_no_substrate,s2f_no_force
 ```
 
 A single run: `python -m experiments.train --variant unet --seed 0`.
@@ -66,22 +66,22 @@ A single run: `python -m experiments.train --variant unet --seed 0`.
 - A model is never split across GPUs: with batch size 2, that would harm batch-norm/GAN training.
 - `train.py` refuses to run without CUDA unless `--allow-cpu` is given.
 - Augmentation runs in `--num-workers` DataLoader processes (default 2 per job; use 0 if RAM is short).
-- Finished runs are skipped and interrupted runs resume from `last_checkpoint.pth`, so the command can simply be re-run.
-- Run time on an 11 GB GPU (2080 Ti class) is about 4-5 min per epoch, i.e. about 7-8 h per variant (100 epochs). On two GPUs:
-  - the 5 core variants take about 16-24 h;
-  - all 14 variants take about 55 h.
-- Progress is logged to `experiments/runs/<variant>_seed<seed>/train.log` and `history.json`.
+- Finished runs (and finished phases) are skipped and interrupted runs resume from the phase's `last_checkpoint.pth`, so the command can simply be re-run.
+- Run time on an RTX A6000: S2F about 3 min per epoch, i.e. about 7 h per variant (100 + 40 epochs); U-Net about 2.5 h. On two GPUs the 6 core variants take about 15 h (`s2f_no_force` reuses phase 1 of `s2f_full`).
+- Progress is logged to `experiments/runs/<variant>_seed<seed>/train.log` (validation pixel r and WFM-RME every epoch) and `phase*/history.json`.
 
-For a quick check that everything runs, add `--epochs 1 --max-batches 2` (and `--allow-cpu` without a GPU).
+For a quick check that everything runs, add `--epochs 1 --max-batches 2` (1 epoch per phase; add `--allow-cpu` without a GPU).
 
 ## Training protocol (all variants)
 
-Every variant is trained from scratch on `alltune_train` (280 pairs) with the same settings:
-- 100 epochs, batch size 2, 1024 × 1024 px, the augmentations of `data/augmentations.py`;
-- Adam (lr 1e-4, β = 0.5/0.999), cosine decay after 10 % linear warm-up;
-- loss = 10 × per-image WFM-RME + 1 × adversarial (PatchGAN, BCE) + 100 × MSE of mean intensity.
+Every variant is trained from scratch on `alltune_train` (280 pairs), batch size 2, 1024 × 1024 px, the augmentations of `data/augmentations.py`, Adam (β = 0.5/0.999) with cosine decay after 10 % linear warm-up, in two phases that mirror the staged training of the released model:
 
-The checkpoint with the lowest reconstruction loss on `alltune_val` (81 pairs) is kept. Variants change one component each (`configs.py`):
+| phase | epochs | reconstruction loss | other terms | lr | starts from |
+|---|---|---|---|---|---|
+| 1 localisation | 100 | 100 × weighted MSE + soft Dice (pos_weight 3, threshold 0.2, Dice 0.5) | 1 × adversarial | 1e-4 | random weights |
+| 2 magnitude | 40 | 10 × per-image WFM-RME | 1 × adversarial + 100 × MSE of mean intensity | 1e-5 | best phase-1 generator |
+
+Each phase uses a new discriminator and optimiser and keeps the checkpoint with the lowest validation loss of that phase on `alltune_val` (81 pairs); the final model is the best phase-2 checkpoint. WFM-RME alone does not learn to localise the forces from random weights (a first protocol that used it from the start gave pixel r ≈ 0), which is also why the released model was first trained with the weighted MSE + Dice loss. Variants change one component each (`configs.py`):
 
 | variant | change |
 |---|---|
@@ -90,7 +90,7 @@ The checkpoint with the lowest reconstruction loss on `alltune_val` (81 pairs) i
 | `unet_gan` | plain U-Net + PatchGAN (pix2pix-style) |
 | `mask_only` | binary cell mask instead of bright field (needs `cell_mask.tif` in each sample folder) |
 | `s2f_no_gan` | no discriminator (λ_GAN = 0) |
-| `s2f_no_force` | no stress-consistency term (λ_force = 0) |
+| `s2f_no_force` | no stress-consistency term (λ_force = 0 in phase 2; phase 1 is identical to `s2f_full` and reused) |
 | `s2f_gan_0.1`, `s2f_gan_5` | λ_GAN = 0.1 / 5 |
 | `s2f_no_substrate` | bright field only (no substrate channels) |
 | `s2f_nominal_substrate` | substrate channels from the nominal Young's modulus of each image |
@@ -99,7 +99,7 @@ The checkpoint with the lowest reconstruction loss on `alltune_val` (81 pairs) i
 | `s2f_bridge_plain` | bottleneck without CBAM, Global Context Module and hierarchical attention |
 | `s2f_no_attention_gates` | skip connections without attention gates |
 
-The released checkpoint was obtained by staged fine-tuning (see Methods). `evaluate.py` reports it alongside the from-scratch variants.
+The released checkpoint was obtained by staged fine-tuning from earlier checkpoints, with substrate channels added only in its last two stages (see Methods). `evaluate.py` reports it alongside the from-scratch variants.
 
 ## Metrics
 

experiments/configs.py

@@ -1,25 +1,32 @@
 """
 Training protocol shared by every baseline / ablation run, and the variants.
 
-All variants are trained FROM SCRATCH with the same protocol, split, number of epochs and
+All variants are trained FROM SCRATCH with the same two-phase protocol, split, number of epochs and
 checkpoint-selection rule, so that differences are attributable to the ablated component only.
-(The released single-cell checkpoint was obtained by staged fine-tuning, see the Methods; the
-'s2f_full' variant is its from-scratch counterpart under this protocol.)
+The two phases mirror the staged training of the released model:
+  phase 1  localisation:  weighted MSE + soft Dice (lambda_rec 100), lr 1e-4    (as train_pix2pix.ipynb)
+  phase 2  magnitude:     per-image WFM-RME (lambda_rec 10) + force consistency (lambda_force 100),
+                          lr 1e-5, initialised from the best phase-1 generator  (as train_pix2pix_settings_mean.ipynb)
+WFM-RME alone does not learn to localise forces from random weights (tested: pixel r ~ 0 after 100 epochs),
+which is why the released model, too, was first trained with the weighted MSE + Dice loss.
+Unlike the released model, every variant sees its full input (incl. substrate channels) from the start.
 """
 import copy
 
 PROTOCOL = dict(
     train_manifest="alltune_train",      # 280 pairs, released model's training split
     val_manifest="alltune_val",          # 81 pairs, used for checkpoint selection only
-    epochs=100,
+    phases=[
+        dict(name="phase1", epochs=100, rec_loss="wmd", lambda_rec=100.0, lambda_force=0.0, g_lr=1e-4, d_lr=1e-4),
+        dict(name="phase2", epochs=40, rec_loss="rme", lambda_rec=10.0, lambda_force="from_variant", g_lr=1e-5, d_lr=1e-5),
+    ],
     batch_size=2,
     img_size=1024,
     augment=True,                        # data/augmentations.AdvancedAugmentations
-    g_lr=1e-4, d_lr=1e-4, beta1=0.5, beta2=0.999,
-    warmup_frac=0.1,                     # cosine schedule with linear warm-up
-    lambda_rec=10.0,                     # per-image WFM-RME
-    lambda_gan=1.0,                      # vanilla (BCE) PatchGAN loss
-    lambda_force=100.0,                  # MSE between mean predicted and mean target intensity
+    beta1=0.5, beta2=0.999,
+    warmup_frac=0.1,                     # cosine schedule with linear warm-up (per phase)
+    lambda_gan=1.0,                      # vanilla (BCE) PatchGAN loss, both phases
+    lambda_force=100.0,                  # phase 2: MSE between mean predicted and mean target intensity
     arch="s2f", in_channels=3, use_multi_scale_input=True, bridge="cbam", attention_gates=True,
     substrate_mode="legacy",             # see experiments/manifest_data.py
     input_mode="bf",
@@ -35,7 +42,7 @@ VARIANTS = {
     "mask_only":              dict(input_mode="mask"),                    # binary cell mask instead of bright field
     # loss ablations
     "s2f_no_gan":             dict(lambda_gan=0.0),
-    "s2f_no_force":           dict(lambda_force=0.0),
+    "s2f_no_force":           dict(lambda_force=0.0, share_phase1_with="s2f_full"),   # force term is a phase-2 term only
     "s2f_gan_0.1":            dict(lambda_gan=0.1),
     "s2f_gan_5":              dict(lambda_gan=5.0),
     # input ablations
@@ -48,13 +55,15 @@ VARIANTS = {
     "s2f_no_attention_gates": dict(attention_gates=False),
 }
 
-# Recommended order if GPU time is limited (about 7-8 h per run of 100 epochs on an 11 GB GPU)
-PRIORITY = ["s2f_full", "unet", "s2f_no_gan", "mask_only", "s2f_no_substrate",
-            "s2f_no_multiscale", "s2f_bridge_plain", "s2f_no_force", "s2f_no_attention_gates",
-            "s2f_bridge_no_cbam", "s2f_nominal_substrate", "unet_gan", "s2f_gan_0.1", "s2f_gan_5"]
+# Recommended order if GPU time is limited (S2F about 7 h per run on an RTX A6000, U-Net about 2.5 h)
+PRIORITY = ["s2f_full", "unet", "unet_gan", "s2f_no_gan", "s2f_no_substrate", "s2f_no_force",
+            "s2f_no_attention_gates", "s2f_no_multiscale", "s2f_bridge_plain",
+            "s2f_bridge_no_cbam", "s2f_nominal_substrate", "s2f_gan_0.1", "s2f_gan_5"]
+# 'mask_only' needs a cell mask (cell_mask.tif) in every sample folder and is therefore not in PRIORITY.
 
 
-def get_config(variant, **overrides):
+def get_config(variant, epochs=None, g_lr=None, d_lr=None, **overrides):
+    """epochs / g_lr / d_lr override every phase (smoke tests); other overrides are top-level keys."""
     if variant not in VARIANTS:
         raise KeyError(f"unknown variant '{variant}'. Options: {', '.join(VARIANTS)}")
     cfg = copy.deepcopy(PROTOCOL)
@@ -63,4 +72,11 @@ def get_config(variant, **overrides):
     cfg["variant"] = variant
     if cfg["substrate_mode"] == "none":
         cfg["in_channels"] = 1
+    for ph in cfg["phases"]:
+        if ph["lambda_force"] == "from_variant":
+            ph["lambda_force"] = cfg["lambda_force"]
+        for k, v in (("epochs", epochs), ("g_lr", g_lr), ("d_lr", d_lr)):
+            if v is not None:
+                ph[k] = v
+    cfg["epochs"] = sum(ph["epochs"] for ph in cfg["phases"])     # total, used to detect finished runs
     return cfg

experiments/losses.py

@@ -21,3 +21,24 @@ class PerImageWFMRMELoss(nn.Module):
         rel = (p - t).abs() / (t + self.eps)
         w = t / (t.mean(dim=(1, 2), keepdim=True) + self.eps)
         return (rel * w).mean()
+
+
+class WeightedMSEPlusDiceLoss(nn.Module):
+    """Loss of the image-only training stage of the released model (Archive/src/utils/losses.py,
+    used in Archive/notebooks/train_pix2pix.ipynb with pos_weight=3, threshold=0.2, dice_weight=0.5).
+    Pixels above `threshold` get weight 1 + pos_weight in the MSE; the soft-Dice term rewards overlap of
+    the predicted and true high-force regions. Unlike L1 / WFM-RME it drives the network to localise
+    the force-generating regions when training starts from random weights."""
+    def __init__(self, pos_weight=3.0, threshold=0.2, dice_weight=0.5):
+        super().__init__()
+        self.pos_weight, self.threshold, self.dice_weight = pos_weight, threshold, dice_weight
+
+    def forward(self, pred, target):
+        w = 1.0 + self.pos_weight * (target > self.threshold).float()
+        mse = (w * (pred - target) ** 2).mean()
+        p, t = pred.flatten(1), target.flatten(1)
+        dice = (2 * (p * t).sum(1) + 1e-6) / (p.sum(1) + t.sum(1) + 1e-6)
+        return mse + self.dice_weight * (1 - dice.mean())
+
+
+REC_LOSSES = {"rme": PerImageWFMRMELoss, "wmd": WeightedMSEPlusDiceLoss}

experiments/train.py

@@ -3,15 +3,16 @@ Train one baseline / ablation variant with the shared protocol (experiments/conf
 
     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 from last_checkpoint.pth
-    python -m experiments.train --variant unet --epochs 1 --max-batches 5   # smoke test
+    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)
 
-Outputs in experiments/runs/<variant>_seed<seed>/:
+Two phases (see experiments/configs.py). Outputs in experiments/runs/<variant>_seed<seed>/:
     config.json            exact configuration used
-    history.json           per-epoch losses (each generator loss term separately), lr, time
-    best_checkpoint.pth    generator weights at the lowest validation reconstruction loss
-    last_checkpoint.pth    full state for --resume
-    visualizations/        predictions on the first validation batch every 10 epochs
+    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.
 """
@@ -20,6 +21,7 @@ import json
 import math
 import os
 import random
+import shutil
 import time
 
 import numpy as np
@@ -30,7 +32,7 @@ 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 PerImageWFMRMELoss
+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
@@ -77,60 +79,45 @@ def save_vis(G, val_loader, cfg, norm, device, path):
     fig.tight_layout(); fig.savefig(path, dpi=100); plt.close(fig)
 
 
-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(os.path.join(run_dir, "visualizations"), 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:                      # check that worker processes work on this machine
-        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)
+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
-    G = build_generator(**cfg).to(device)
     D = build_discriminator(cfg["in_channels"]).to(device) if use_gan else None
-    cfg = dict(cfg, n_params_generator=n_params(G), n_params_discriminator=n_params(D) if D else 0,
-               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 {cfg['n_params_discriminator']/1e6:.2f} M; "
-          f"train {cfg['n_train']}, val {cfg['n_val']}, device {device}")
-
-    g_opt = torch.optim.Adam(G.parameters(), lr=cfg["g_lr"], betas=(cfg["beta1"], cfg["beta2"]))
-    d_opt = torch.optim.Adam(D.parameters(), lr=cfg["d_lr"], betas=(cfg["beta1"], cfg["beta2"])) if use_gan else None
-    steps = len(train_loader) * cfg["epochs"]
+    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, bce = PerImageWFMRMELoss(), nn.BCEWithLogitsLoss()
+    rec_loss, rme, bce = REC_LOSSES[ph["rec_loss"]](), PerImageWFMRMELoss(), nn.BCEWithLogitsLoss()
     hist, best, start = [], float("inf"), 0
-
-    last_path = os.path.join(run_dir, "last_checkpoint.pth")
+    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"resumed at epoch {start}")
+        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, cfg["epochs"]):
+    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):
@@ -144,7 +131,7 @@ def train(cfg, run_dir, dataset_root=DEFAULT_DATASET_ROOT, device=None, resume=F
             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 = cfg["lambda_rec"] * l_rec + cfg["lambda_gan"] * l_gan + cfg["lambda_force"] * l_force
+            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:
@@ -153,37 +140,98 @@ def train(cfg, run_dir, dataset_root=DEFAULT_DATASET_ROOT, device=None, resume=F
                 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)
+                acc[k] += float(v.detach())
             nb += 1
-        G.eval(); v_rec = v_mse = 0.0; nv = 0
+        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_mse += float(F.mse_loss(p, y)); nv += 1
-        h = {"epoch": epoch, **{f"train_{k}": v / max(nb, 1) for k, v in acc.items()},
-             "val_rec": v_rec / max(nv, 1), "val_mse": v_mse / max(nv, 1),
+                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"epoch {epoch:3d}  G {h['train_g']:.4f} (rec {h['train_rec']:.4f} gan {h['train_gan']:.4f} "
+        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"{'*' if is_best else ' '}  {h['epoch_seconds']:.0f}s", flush=True)
-        json.dump(hist, open(os.path.join(run_dir, "history.json"), "w"), indent=1)
+              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, "generator_state_dict": G.state_dict(), "config": cfg,
-                        "val_loss": h["val_rec"], "history": hist}, os.path.join(run_dir, "best_checkpoint.pth"))
+            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)
-        if epoch % 10 == 0 or epoch == cfg["epochs"] - 1:
-            save_vis(G, val_loader, cfg, norm, device, os.path.join(run_dir, "visualizations", f"epoch_{epoch:03d}.png"))
-    print(f"done. best val_rec {best:.4f}; checkpoints in {run_dir}")
+        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:                      # check that worker processes work on this machine
+        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"])          # next phase starts from this phase's best
+    best_ck["config"] = cfg
+    torch.save(best_ck, os.path.join(run_dir, "best_checkpoint.pth"))   # final model = best of the last phase
+    json.dump(all_hist, open(os.path.join(run_dir, "history.json"), "w"), indent=1)   # written last = run finished
+    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