Spaces:
Running on CPU Upgrade
Running on CPU Upgrade
| """Retrain the deployed ReStraV checkpoint with the corrected (official) recipe. | |
| The previous restrav.pth was trained on ImageNet-normalized frames but is *served* raw [0,1] | |
| (registry video_norm = identity) -- a train/serve mismatch that collapsed it (AUROC ~0.50, | |
| real_acc ~2%). This retrains the SynthScope `SynthScopeReStraVModel` head end-to-end on raw [0,1] frames, | |
| so training and serving use the exact same DINOv2 features. | |
| Key fidelity choices: | |
| * Features come from the model's OWN ``_encode`` + ``_geometry_features`` (so they are | |
| bit-identical to what the deployed forward computes, incl. the official [mean,min,max,var] | |
| summary order), with frames decoded by the SAME serving pipeline (media.video_clip_tensor, | |
| center window + BICUBIC, raw [0,1]). | |
| * Official training recipe: standardize (BatchNorm1d running stats, train mode), MLP | |
| 21->64->32->1, BCEWithLogits, plain Adam lr 1e-3 (no weight decay), 10 epochs, batch 32. | |
| * Save the FINAL-epoch state_dict (no val-loss checkpoint selection). | |
| Run: conda run -n SynthScope python interface/scripts/retrain_restrav.py [--n 3000] | |
| """ | |
| import argparse | |
| import glob | |
| import random | |
| import sys | |
| import time | |
| from pathlib import Path | |
| import torch | |
| import torch.nn as nn | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[1])) # interface/ | |
| sys.path.insert(0, str(Path(__file__).resolve().parents[2])) # repo root (config_utils) | |
| from backend.registry import get_spec # noqa: E402 | |
| from backend import media # noqa: E402 | |
| from backend.engine import _build_model # noqa: E402 | |
| DEV = "cuda" if torch.cuda.is_available() else "cpu" | |
| TRAIN_ROOT = "/hddstore/ssalehi/SynthScope/restrav_data/train/genvideo" | |
| OUT_DIR = Path("/hddstore/ssalehi/SynthScope/checkpoints/_train/restrav_fixed") | |
| DEPLOY_LINK = Path("/hddstore/ssalehi/SynthScope/checkpoints/restrav.pth") | |
| def build_model(spec): | |
| from config_utils import compose_config | |
| cfg = compose_config(spec.overrides) | |
| model = _build_model(cfg).to(DEV).float() | |
| model.eval() # frozen backbone stays in eval | |
| return model | |
| def extract_features(model, spec, items, tag): | |
| X, y = [], [] | |
| t0 = time.time() | |
| for i, (p, lab) in enumerate(items): | |
| try: | |
| clip = media.video_clip_tensor( | |
| p, spec.num_frames, spec.fps, spec.frame_size, | |
| window_center=spec.window_center, interp=spec.resize_interp, | |
| ) # [T,3,H,W] raw [0,1] | |
| z = model._encode(clip.unsqueeze(0).to(DEV)) # [1,T,D] | |
| feat = model._geometry_features(z)[0] # [21] | |
| if torch.isfinite(feat).all(): | |
| X.append(feat.cpu()); y.append(lab) | |
| except Exception as e: | |
| if i < 5: | |
| print(f" skip {Path(p).name}: {e}", flush=True) | |
| if (i + 1) % 300 == 0: | |
| print(f" [{tag}] {i+1}/{len(items)} ({(time.time()-t0)/(i+1):.2f}s/vid)", flush=True) | |
| return torch.stack(X), torch.tensor(y, dtype=torch.float32) | |
| def main(): | |
| ap = argparse.ArgumentParser() | |
| ap.add_argument("--n", type=int, default=3000, help="videos per class") | |
| ap.add_argument("--epochs", type=int, default=10) | |
| args = ap.parse_args() | |
| spec = get_spec("restrav") | |
| model = build_model(spec) | |
| print(f"model built on {DEV}; serving recipe: num_frames={spec.num_frames} " | |
| f"window_center={spec.window_center} interp={spec.resize_interp} video_norm={spec.video_norm}", flush=True) | |
| reals = sorted(glob.glob(TRAIN_ROOT + "/0_real/*.mp4")) | |
| fakes = sorted(glob.glob(TRAIN_ROOT + "/1_fake/*.mp4")) | |
| random.Random(0).shuffle(reals); random.Random(1).shuffle(fakes) | |
| # SynthScope label convention: 0 = real, 1 = fake. | |
| items = [(p, 0) for p in reals[:args.n]] + [(p, 1) for p in fakes[:args.n]] | |
| random.Random(2).shuffle(items) | |
| print(f"extracting features for {len(items)} videos (raw [0,1], model._encode)...", flush=True) | |
| X, y = extract_features(model, spec, items, "train") | |
| X = X.to(DEV); y = y.to(DEV) | |
| print(f"features: {tuple(X.shape)} reals={(y==0).sum().item()} fakes={(y==1).sum().item()}", flush=True) | |
| # --- train head (official recipe). feature_norm is BatchNorm1d(affine=False): no params, | |
| # just running stats accumulated in train mode. classifier holds the trainable weights. --- | |
| model.feature_norm.train(); model.classifier.train() | |
| opt = torch.optim.Adam(model.classifier.parameters(), lr=1e-3) # wd=0 (plain Adam) | |
| crit = nn.BCEWithLogitsLoss() | |
| for ep in range(args.epochs): | |
| perm = torch.randperm(len(X), device=DEV) | |
| tot = 0.0 | |
| for i in range(0, len(X), 32): | |
| b = perm[i:i+32] | |
| opt.zero_grad() | |
| logit = model.classifier(model.feature_norm(X[b])).squeeze(-1) # fake logit | |
| loss = crit(logit, y[b]) | |
| loss.backward(); opt.step() | |
| tot += loss.item() * len(b) | |
| print(f" epoch {ep+1}/{args.epochs} loss={tot/len(X):.4f}", flush=True) | |
| # --- save FINAL state_dict (matches what engine._load_checkpoint expects) --- | |
| model.eval() | |
| OUT_DIR.mkdir(parents=True, exist_ok=True) | |
| out = OUT_DIR / "model_final.pth" | |
| torch.save(model.state_dict(), out) | |
| print(f"saved {out}", flush=True) | |
| # quick in-sample sanity (real_acc / fake_acc on the train features) | |
| with torch.no_grad(): | |
| p_fake = torch.sigmoid(model.classifier(model.feature_norm(X)).squeeze(-1)) | |
| pred = (p_fake > 0.5).float() | |
| racc = (pred[y == 0] == 0).float().mean().item() | |
| facc = (pred[y == 1] == 1).float().mean().item() | |
| print(f"in-sample: real_acc={racc*100:.1f} fake_acc={facc*100:.1f}", flush=True) | |
| # relink the deployed checkpoint | |
| if DEPLOY_LINK.is_symlink() or DEPLOY_LINK.exists(): | |
| DEPLOY_LINK.unlink() | |
| DEPLOY_LINK.symlink_to(out) | |
| print(f"relinked {DEPLOY_LINK} -> {out}", flush=True) | |
| if __name__ == "__main__": | |
| main() | |