| """Evaluation script for few-shot RS-IMLE models. |
| |
| Computes FID (5000 samples) and Precision/Recall (1000 samples) at one or |
| more latent noise scales, and saves a sample grid. |
| """ |
|
|
| import argparse |
| import json |
| import os |
| import sys |
| import shutil |
| import time |
|
|
| sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
|
|
| import imageio |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| import torchvision |
| from distutils.util import strtobool |
|
|
| from data import set_up_data |
| from hps import Hyperparams, add_imle_arguments, parse_args_and_update_hparams |
| from models import IMLE |
| from sampler import Sampler |
|
|
| from cleanfid import fid |
| import cleanfid.features as _cleanfid_feat |
| import cleanfid.inception_torchscript as _cleanfid_incept |
|
|
| _inception_cache = os.path.join(os.path.expanduser("~"), ".cache", "cleanfid") |
| os.makedirs(_inception_cache, exist_ok=True) |
| _orig_feature_extractor = _cleanfid_feat.feature_extractor |
|
|
|
|
| def _patched_feature_extractor(name="torchscript_inception", |
| device=torch.device("cuda"), |
| resize_inside=False, use_dataparallel=True): |
| if name == "torchscript_inception": |
| model = _cleanfid_incept.InceptionV3W( |
| _inception_cache, download=True, resize_inside=resize_inside |
| ).to(device) |
| model.eval() |
| if use_dataparallel: |
| model = torch.nn.DataParallel(model) |
| return lambda x: model(x) |
| return _orig_feature_extractor(name, device, resize_inside, use_dataparallel) |
|
|
|
|
| _cleanfid_feat.feature_extractor = _patched_feature_extractor |
|
|
|
|
| |
| |
| |
|
|
| def generate_images_to_dir(model, latent_dim, num_images, out_dir, |
| batch_size=16, noise_scale=1.0): |
| """Generate images with latents sampled as z ~ N(0, noise_scale^2).""" |
| os.makedirs(out_dir, exist_ok=True) |
| device = next(model.parameters()).device |
| idx = 0 |
| with torch.no_grad(): |
| while idx < num_images: |
| bs = min(batch_size, num_images - idx) |
| z = torch.randn(bs, latent_dim, device=device) * noise_scale |
| imgs = model(z, None) |
| imgs = (imgs + 1.0) * 127.5 |
| imgs = imgs.clamp(0, 255).permute(0, 2, 3, 1) |
| imgs = imgs.cpu().numpy().astype(np.uint8) |
| for j in range(bs): |
| imageio.imwrite(os.path.join(out_dir, f"{idx}.png"), imgs[j]) |
| idx += 1 |
|
|
|
|
| def generate_grid_image(model, latent_dim, num_samples, nrow, out_path, |
| noise_scale=1.0): |
| device = next(model.parameters()).device |
| with torch.no_grad(): |
| z = torch.randn(num_samples, latent_dim, device=device) * noise_scale |
| imgs = model(z, None) |
| imgs = (imgs + 1.0) / 2.0 |
| imgs = imgs.clamp(0.0, 1.0) |
| grid = torchvision.utils.make_grid(imgs, nrow=nrow, padding=2) |
| grid_pil = torchvision.transforms.functional.to_pil_image(grid.cpu()) |
| grid_pil.save(out_path) |
| print(f" Saved grid ({num_samples} images) to {out_path}") |
|
|
|
|
| def _slerp_train_style(a: torch.Tensor, b: torch.Tensor, t: torch.Tensor) -> torch.Tensor: |
| """Same spherical interpolation as train.py (FID-epoch interp strips).""" |
| a = F.normalize(a, dim=-1) |
| b = F.normalize(b, dim=-1) |
| dot = torch.sum(a * b, dim=-1, keepdim=True).clamp(-1.0, 1.0) |
| omega = torch.acos(dot) |
| sin_omega = torch.sin(omega) |
| t = t.view(-1, 1) |
| factor1 = torch.sin((1.0 - t) * omega) / sin_omega |
| factor2 = torch.sin(t * omega) / sin_omega |
| return factor1 * a + factor2 * b |
|
|
|
|
| def generate_slerp_grid(model, H, sampler, interp_pairs, interp_steps, nrow, |
| out_path, noise_scale=1.0): |
| """SLERP strips like train.py / supplementary few-shot figure (latent geodesics).""" |
| device = next(model.parameters()).device |
| with torch.no_grad(): |
| z1 = torch.randn(interp_pairs, H.latent_dim, device=device, |
| dtype=torch.float32) * noise_scale |
| z2 = torch.randn(interp_pairs, H.latent_dim, device=device, |
| dtype=torch.float32) * noise_scale |
| t_vals = torch.linspace(0.0, 1.0, interp_steps, device=device, |
| dtype=torch.float32) |
| all_rows = [] |
| for pi in range(interp_pairs): |
| z_interp = _slerp_train_style( |
| z1[pi:pi + 1].repeat(interp_steps, 1), |
| z2[pi:pi + 1].repeat(interp_steps, 1), |
| t_vals, |
| ) |
| snoise_tmp = [s[:interp_steps].normal_() for s in sampler.snoise_tmp] |
| preds = sampler.sample(z_interp, model, snoise_tmp) |
| preds_t = torch.from_numpy(preds).float() / 255.0 |
| preds_t = preds_t.permute(0, 3, 1, 2) |
| all_rows.append(preds_t) |
| tensor = torch.cat(all_rows, dim=0) |
| grid = torchvision.utils.make_grid(tensor, nrow=nrow, padding=2) |
| grid_pil = torchvision.transforms.functional.to_pil_image( |
| grid.cpu().clamp(0.0, 1.0)) |
| grid_pil.save(out_path) |
| n = interp_pairs * interp_steps |
| print(f" Saved SLERP grid ({n} images, {interp_pairs}x{interp_steps}) to {out_path}") |
|
|
|
|
| |
| |
| |
|
|
| def override_model_cycles(model, test_H_cycles=None, test_L_cycles=None, |
| test_refinement_steps=None): |
| raw = model.module if hasattr(model, 'module') else model |
| mapper = raw.decoder.mapping_network |
|
|
| inner = getattr(mapper, 'trm', None) |
| if inner is None: |
| print(" No TRM inner module found; cycle override skipped.") |
| return |
|
|
| old_H = inner.H_cycles |
| old_L = inner.L_cycles |
| old_halt = mapper.refinement_steps |
|
|
| if test_H_cycles is not None: |
| inner.H_cycles = test_H_cycles |
| if test_L_cycles is not None: |
| inner.L_cycles = test_L_cycles |
| if test_refinement_steps is not None: |
| mapper.refinement_steps = test_refinement_steps |
|
|
| print(f" Cycle override: H {old_H}->{inner.H_cycles}, " |
| f"L {old_L}->{inner.L_cycles}, refinement {old_halt}->{mapper.refinement_steps}") |
|
|
|
|
| def build_and_load_model(H): |
| """Match helpers.train_helpers.load_imle: strip ``module.``, then DP main only.""" |
| from helpers.train_helpers import restore_params |
|
|
| local_rank = getattr(H, 'local_rank', 0) |
| mpi_size = getattr(H, 'mpi_size', 1) |
| strict = bool(H.load_strict) |
|
|
| model = IMLE(H) |
| if H.restore_path and os.path.isfile(H.restore_path): |
| print(f" Loading model: {H.restore_path}") |
| restore_params( |
| model, H.restore_path, local_rank, mpi_size, |
| map_cpu=True, strict=strict, |
| ) |
| else: |
| print(" WARNING: restore_path not found or not set, using random weights!") |
|
|
| model = torch.nn.DataParallel(model.cuda()) |
|
|
| ema_model = None |
| if getattr(H, 'restore_ema_path', None) and os.path.isfile(H.restore_ema_path): |
| print(f" Loading EMA: {H.restore_ema_path}") |
| ema_model = IMLE(H) |
| restore_params( |
| ema_model, H.restore_ema_path, local_rank, mpi_size, |
| map_cpu=True, strict=strict, |
| ) |
| ema_model = ema_model.cuda() |
| ema_model.requires_grad_(False) |
|
|
| return model, ema_model |
|
|
|
|
| |
| |
| |
|
|
| def evaluate_at_scale(model, H, noise_scale, output_dir, batch_size, |
| num_fid, num_pr, num_grid, grid_nrow, sampler=None, |
| grid_mode='iid'): |
| results = {'noise_scale': noise_scale} |
|
|
| if num_grid > 0: |
| grid_path = os.path.join(output_dir, f"grid_scale_{noise_scale:.2f}.png") |
| mode = (grid_mode or 'iid').lower() |
| if mode == 'slerp': |
| if sampler is None: |
| raise ValueError( |
| "grid_mode=slerp requires a Sampler (set_up_data + Sampler(...))") |
| if grid_nrow <= 0 or num_grid % grid_nrow != 0: |
| raise ValueError( |
| f"For slerp, num_grid_samples ({num_grid}) must be divisible by " |
| f"grid_nrow ({grid_nrow}); each row is one SLERP chain.") |
| interp_steps = grid_nrow |
| interp_pairs = num_grid // interp_steps |
| generate_slerp_grid( |
| model, H, sampler, interp_pairs, interp_steps, interp_steps, |
| grid_path, noise_scale=noise_scale) |
| else: |
| generate_grid_image(model, H.latent_dim, num_grid, grid_nrow, grid_path, |
| noise_scale=noise_scale) |
| results['grid_path'] = grid_path |
|
|
| ref_dir = f'{H.data_root}/img' |
|
|
| |
| if num_fid > 0: |
| fid_dir = os.path.join(output_dir, f"_tmp_fid_{noise_scale:.2f}") |
| os.makedirs(fid_dir, exist_ok=True) |
| t0 = time.time() |
| generate_images_to_dir(model, H.latent_dim, num_fid, fid_dir, |
| batch_size=batch_size, noise_scale=noise_scale) |
| cur_fid = fid.compute_fid(ref_dir, fid_dir, verbose=False, num_workers=0) |
| results['fid'] = cur_fid |
| print(f" FID = {cur_fid:.4f} ({time.time() - t0:.1f}s)") |
| shutil.rmtree(fid_dir, ignore_errors=True) |
|
|
| |
| if num_pr > 0: |
| pr_dir = os.path.join(output_dir, f"_tmp_pr_{noise_scale:.2f}") |
| os.makedirs(pr_dir, exist_ok=True) |
| t0 = time.time() |
| generate_images_to_dir(model, H.latent_dim, num_pr, pr_dir, |
| batch_size=batch_size, noise_scale=noise_scale) |
| try: |
| from helpers.improved_precision_recall import compute_prec_recall |
| precision, recall = compute_prec_recall(ref_dir, pr_dir) |
| results['precision'] = precision |
| results['recall'] = recall |
| print(f" Precision = {precision:.4f}, Recall = {recall:.4f} " |
| f"({time.time() - t0:.1f}s)") |
| except ImportError: |
| print(" WARNING: improved_precision_recall not found, skipping P/R") |
| shutil.rmtree(pr_dir, ignore_errors=True) |
|
|
| return results |
|
|
|
|
| |
| |
| |
|
|
| def main(): |
| H = Hyperparams() |
| parser = argparse.ArgumentParser( |
| description="Fewshot RS-IMLE: noise-resilience evaluation", |
| formatter_class=argparse.RawDescriptionHelpFormatter, |
| epilog=__doc__, |
| ) |
| parser = add_imle_arguments(parser) |
|
|
| parser.add_argument('--noise_scales', type=str, default='1.0,1.5,2.0,2.5,3.0', |
| help='Comma-separated noise scale factors for latent z ' |
| '(1.0 = standard Gaussian, >1 = heavier tails)') |
| parser.add_argument('--test_H_cycles', type=int, default=None) |
| parser.add_argument('--test_L_cycles', type=int, default=None) |
| parser.add_argument('--test_refinement_steps', type=int, default=None) |
| |
| parser.add_argument('--num_pr_samples', type=int, default=1000) |
| parser.add_argument('--num_grid_samples', type=int, default=40) |
| parser.add_argument('--grid_nrow', type=int, default=8) |
| parser.add_argument('--output_dir', type=str, default=None) |
| parser.add_argument('--use_ema', default=False, |
| type=lambda x: bool(strtobool(x))) |
| parser.add_argument('--test_batch_size', type=int, default=16) |
| parser.add_argument('--grid_only', action='store_true', |
| help='Skip FID and Precision/Recall; only write sample grids. ' |
| 'Sets num_fid_samples and num_pr_samples to 0 and ' |
| 'noise_scales to 1.0.') |
| parser.add_argument('--grid_mode', type=str, default='iid', |
| choices=['iid', 'slerp'], |
| help='iid: independent z samples (default). ' |
| 'slerp: latent SLERP strips as in train.py / ' |
| 'supplementary few-shot figure (needs num_grid_samples ' |
| '= K * grid_nrow for K chains of length grid_nrow).') |
|
|
| parse_args_and_update_hparams(H, parser) |
|
|
| if H.grid_only: |
| H.num_fid_samples = 0 |
| H.num_pr_samples = 0 |
| H.noise_scales = '1.0' |
|
|
| noise_scales = [float(s.strip()) for s in H.noise_scales.split(',')] |
|
|
| if H.output_dir: |
| output_dir = H.output_dir |
| else: |
| output_dir = os.path.join(H.save_dir, 'noise_resilience') |
| os.makedirs(output_dir, exist_ok=True) |
|
|
| H, data_train, data_valid, preprocess_fn = set_up_data(H) |
|
|
| print("=" * 70) |
| if H.grid_only: |
| print(" Fewshot RS-IMLE — Quality grid only (no FID / P/R)") |
| print(f" grid_mode: {getattr(H, 'grid_mode', 'iid')}") |
| else: |
| print(" Fewshot RS-IMLE — Noise Resilience Evaluation") |
| print("=" * 70) |
| print(f" Dataset: {H.dataset}") |
| print(f" Data root: {H.data_root}") |
| print(f" Checkpoint: {H.restore_path}") |
| print(f" EMA checkpoint: {getattr(H, 'restore_ema_path', None)}") |
| print(f" use_ema: {H.use_ema}") |
| print(f" use_rtm: {getattr(H, 'use_rtm', False)}") |
| print(f" latent_dim: {H.latent_dim}") |
| print(f" H/L/refinement: {H.H_cycles}/{H.L_cycles}/{H.refinement_steps}") |
| if H.test_H_cycles is not None or H.test_L_cycles is not None: |
| print(f" test override H/L/refinement: {H.test_H_cycles}/{H.test_L_cycles}/{H.test_refinement_steps}") |
| print(f" Noise scales: {noise_scales}") |
| print(f" FID samples: {H.num_fid_samples}") |
| print(f" P/R samples: {H.num_pr_samples}") |
| print(f" Output dir: {output_dir}") |
| print("=" * 70) |
|
|
| model, ema_model = build_and_load_model(H) |
|
|
| if H.use_ema and ema_model is not None: |
| eval_model = ema_model |
| print(" Using EMA model for evaluation") |
| else: |
| eval_model = model |
| if H.use_ema and ema_model is None: |
| print(" WARNING: --use_ema True but no EMA loaded, using main model") |
| print(" Using main model for evaluation") |
|
|
| override_model_cycles( |
| eval_model, |
| test_H_cycles=H.test_H_cycles, |
| test_L_cycles=H.test_L_cycles, |
| test_refinement_steps=H.test_refinement_steps, |
| ) |
|
|
| eval_model.eval() |
|
|
| sampler = None |
| if H.grid_only and getattr(H, 'grid_mode', 'iid').lower() == 'slerp': |
| sampler = Sampler(H, len(data_train), preprocess_fn) |
| print(f" Sampler dataset len: {len(data_train)} (for IMLE buffers)") |
|
|
| max_samples = max( |
| H.num_fid_samples, |
| H.num_pr_samples, |
| H.num_grid_samples if H.num_grid_samples > 0 else 0, |
| ) |
| if max_samples <= 0: |
| max_samples = max(H.test_batch_size, 16) |
| batch_size = min(H.test_batch_size, max_samples) |
| all_results = [] |
|
|
| for scale in noise_scales: |
| print(f"\n{'─' * 70}") |
| print(f" Noise scale = {scale:.2f} (z ~ N(0, {scale:.2f}²))") |
| print(f"{'─' * 70}") |
| res = evaluate_at_scale( |
| eval_model, H, scale, output_dir, |
| batch_size, H.num_fid_samples, |
| H.num_pr_samples, H.num_grid_samples, |
| H.grid_nrow, |
| sampler=sampler, |
| grid_mode=getattr(H, 'grid_mode', 'iid'), |
| ) |
| all_results.append(res) |
|
|
| |
| raw = eval_model.module if hasattr(eval_model, 'module') else eval_model |
| mapper = raw.decoder.mapping_network |
| inner = getattr(mapper, 'trm', None) |
|
|
| output = { |
| 'config': { |
| 'restore_path': H.restore_path, |
| 'restore_ema_path': getattr(H, 'restore_ema_path', None), |
| 'use_ema': H.use_ema, |
| 'dataset': H.dataset, |
| 'data_root': H.data_root, |
| 'use_rtm': bool(getattr(H, 'use_rtm', False)), |
| 'latent_dim': H.latent_dim, |
| 'H_cycles_trained': H.H_cycles, |
| 'L_cycles_trained': H.L_cycles, |
| 'refinement_steps_trained': H.refinement_steps, |
| 'H_cycles_eval': inner.H_cycles if inner else H.H_cycles, |
| 'L_cycles_eval': inner.L_cycles if inner else H.L_cycles, |
| 'refinement_steps_eval': mapper.refinement_steps if hasattr(mapper, 'refinement_steps') else H.refinement_steps, |
| 'num_fid_samples': H.num_fid_samples, |
| 'num_pr_samples': H.num_pr_samples, |
| 'grid_mode': getattr(H, 'grid_mode', 'iid'), |
| }, |
| 'results': all_results, |
| } |
|
|
| json_path = os.path.join(output_dir, "noise_resilience.json") |
| with open(json_path, "w") as f: |
| json.dump(output, f, indent=2) |
|
|
| |
| summary_lines = [] |
| if H.grid_only: |
| header = f"{'Scale':>7s} {'Grid':>50s}" |
| sep = "─" * len(header) |
| summary_lines.append(sep) |
| summary_lines.append(header) |
| summary_lines.append(sep) |
| for r in all_results: |
| s = r['noise_scale'] |
| g = r.get('grid_path', 'N/A') |
| summary_lines.append(f"{s:>7.2f} {g}") |
| summary_lines.append(sep) |
| else: |
| header = (f"{'Scale':>7s} {'FID':>10s} " |
| f"{'Precision':>10s} {'Recall':>10s}") |
| sep = "─" * len(header) |
| summary_lines.append(sep) |
| summary_lines.append(header) |
| summary_lines.append(sep) |
| for r in all_results: |
| s = r['noise_scale'] |
| f_val = f"{r['fid']:.4f}" if 'fid' in r else "N/A" |
| p_val = f"{r['precision']:.4f}" if 'precision' in r else "N/A" |
| r_val = f"{r['recall']:.4f}" if 'recall' in r else "N/A" |
| summary_lines.append( |
| f"{s:>7.2f} {f_val:>10s} {p_val:>10s} {r_val:>10s}") |
| summary_lines.append(sep) |
| summary_text = "\n".join(summary_lines) |
|
|
| summary_path = os.path.join(output_dir, "summary.txt") |
| with open(summary_path, "w") as f: |
| f.write(summary_text + "\n") |
|
|
| print(f"\n{'=' * 70}") |
| if H.grid_only: |
| print(" QUALITY GRID OUTPUT") |
| else: |
| print(" NOISE RESILIENCE RESULTS") |
| print(f"{'=' * 70}") |
| print(summary_text) |
| print(f"\n Full results: {json_path}") |
| print(f" Summary: {summary_path}") |
| print(f"{'=' * 70}") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
