"""Evaluation script for RS-IMLE (CIFAR-10 and CelebA-HQ-256). Computes FID (5000 samples) and Precision/Recall (1000 samples), and saves a sample grid. Cycle counts (H, L, refinement_steps) can be overridden at test time. """ import argparse import json import os import shutil import time from distutils.util import strtobool import imageio import numpy as np import torch import torchvision from data import set_up_data from hps import Hyperparams, add_imle_arguments, parse_args_and_update_hparams from models import IMLE from torch import autocast # ────────────────────────────────────────────────────────────────────────── # cleanfid monkey-patch # ────────────────────────────────────────────────────────────────────────── 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 # ────────────────────────────────────────────────────────────────────────── # Helpers # ────────────────────────────────────────────────────────────────────────── def generate_images_to_dir(imle_model, latent_dim, num_images, out_dir, batch_size=16): """Generate num_images PNG files into out_dir for FID/P-R evaluation.""" os.makedirs(out_dir, exist_ok=True) device = next(imle_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) with autocast(device_type="cuda"): imgs = imle_model(z, None) # Model output is in [-1, 1], convert to [0, 255] uint8 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(imle_model, latent_dim, num_samples, nrow, out_path): """Generate a grid of num_samples images and save to out_path.""" device = next(imle_model.parameters()).device with torch.no_grad(): z = torch.randn(num_samples, latent_dim, device=device) with autocast(device_type="cuda"): imgs = imle_model(z, None) # Convert [-1, 1] -> [0, 1] for torchvision grid 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 sample grid ({num_samples} images, {nrow} per row) to {out_path}") def override_model_cycles(imle_model, test_H_cycles=None, test_L_cycles=None, test_refinement_steps=None): """Override H_cycles / L_cycles / refinement_steps inside the loaded model.""" model = imle_model.module if hasattr(imle_model, 'module') else imle_model mapper = model.decoder.mapping_network inner = getattr(mapper, 'trm', None) if inner is None: print("WARNING: 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_cycles {old_H}->{inner.H_cycles}, " f"L_cycles {old_L}->{inner.L_cycles}, " f"refinement_steps {old_halt}->{mapper.refinement_steps}") def build_and_load_model(H): """Build IMLE model and load checkpoint weights. Returns (imle_model, ema_model_or_None). """ imle_model = IMLE(H).cuda() ema_model = None if H.restore_path and os.path.isfile(H.restore_path): print(f"Loading model weights from: {H.restore_path}") state_dict = torch.load(H.restore_path, map_location='cuda') # Handle DataParallel-saved checkpoints (keys start with 'module.') has_module_prefix = any(k.startswith('module.') for k in state_dict.keys()) if has_module_prefix: # Wrap model in DataParallel first, then load imle_model = torch.nn.DataParallel(imle_model) imle_model.load_state_dict(state_dict, strict=bool(H.load_strict)) else: imle_model.load_state_dict(state_dict, strict=bool(H.load_strict)) imle_model = torch.nn.DataParallel(imle_model) else: imle_model = torch.nn.DataParallel(imle_model) if H.restore_path: print(f"WARNING: restore_path '{H.restore_path}' not found!") else: print("WARNING: No --restore_path specified, using random weights!") # Load EMA if requested if getattr(H, 'restore_ema_path', None) and os.path.isfile(H.restore_ema_path): print(f"Loading EMA weights from: {H.restore_ema_path}") ema_model = IMLE(H).cuda() ema_state = torch.load(H.restore_ema_path, map_location='cuda') ema_model.load_state_dict(ema_state, strict=bool(H.load_strict)) return imle_model, ema_model # ────────────────────────────────────────────────────────────────────────── # Main # ────────────────────────────────────────────────────────────────────────── def main(): # ── Parse arguments ── H = Hyperparams() parser = argparse.ArgumentParser( description="RS-IMLE Evaluation: FID, Precision/Recall, Sample Grid", formatter_class=argparse.RawDescriptionHelpFormatter, epilog=__doc__, ) parser = add_imle_arguments(parser) # Test-specific arguments parser.add_argument('--test_H_cycles', type=int, default=None, help='Override H_cycles at test time (default: use trained value)') parser.add_argument('--test_L_cycles', type=int, default=None, help='Override L_cycles at test time (default: use trained value)') # Note: --test_refinement_steps and --num_fid_samples are registered by # add_imle_arguments() above; only test-time-only knobs are added here. parser.add_argument('--num_pr_samples', type=int, default=1000, help='Number of generated samples for Precision/Recall') parser.add_argument('--num_grid_samples', type=int, default=40, help='Number of samples in the output grid image') parser.add_argument('--grid_nrow', type=int, default=8, help='Number of images per row in grid (default 8, so 8x5=40)') parser.add_argument('--output_dir', type=str, default=None, help='Output directory (default: {save_dir}/test_results)') parser.add_argument('--use_ema', default=False, type=lambda x: bool(strtobool(x)), help='Evaluate EMA model instead of main model') parser.add_argument('--skip_fid', default=False, type=lambda x: bool(strtobool(x)), help='Skip FID computation') parser.add_argument('--skip_pr', default=False, type=lambda x: bool(strtobool(x)), help='Skip Precision/Recall computation') parser.add_argument('--test_batch_size', type=int, default=128, help='Batch size for generating images during test') parser.add_argument('--dump_samples_dir', type=str, default=None, help='If set, dump --dump_samples_n PNGs here. ' 'Combine with --skip_fid True --skip_pr True to skip metrics.') parser.add_argument('--dump_samples_n', type=int, default=256, help='Number of PNGs to save when --dump_samples_dir is set.') parse_args_and_update_hparams(H, parser) # ── Setup output directory ── if H.output_dir: output_dir = H.output_dir else: output_dir = os.path.join(H.save_dir, 'test_results') os.makedirs(output_dir, exist_ok=True) # ── Setup data (needed for FID reference path) ── H, data_train, data_valid, preprocess_fn = set_up_data(H) print("=" * 60) print("RS-IMLE Evaluation") print("=" * 60) print(f" Dataset: {H.dataset}") print(f" Data root: {H.data_root}") print(f" Checkpoint: {H.restore_path}") print(f" use_rtm: {getattr(H, 'use_rtm', False)}") print(f" H_cycles: {H.H_cycles}") print(f" L_cycles: {H.L_cycles}") print(f" refinement_steps: {H.refinement_steps}") if H.test_H_cycles is not None: print(f" test_H_cycles: {H.test_H_cycles} (OVERRIDE)") if H.test_L_cycles is not None: print(f" test_L_cycles: {H.test_L_cycles} (OVERRIDE)") if H.test_refinement_steps is not None: print(f" test_refinement_steps: {H.test_refinement_steps} (OVERRIDE)") print(f" Output dir: {output_dir}") print(f" FID samples: {H.num_fid_samples}") print(f" P/R samples: {H.num_pr_samples}") print(f" Grid samples: {H.num_grid_samples}") print("=" * 60) # ── Build and load model ── imle_model, ema_model = build_and_load_model(H) # Decide which model to evaluate if H.use_ema and ema_model is not None: eval_model = ema_model print("Evaluating EMA model") else: eval_model = imle_model if H.use_ema and ema_model is None: print("WARNING: --use_ema True but no EMA weights loaded, using main model") print("Evaluating main model") # ── Override cycles at test time ── 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() batch_size = min(getattr(H, 'test_batch_size', 128), H.num_fid_samples) results = {} if H.dump_samples_dir: print(f"\nDumping {H.dump_samples_n} PNG samples to {H.dump_samples_dir}") generate_images_to_dir(eval_model, H.latent_dim, H.dump_samples_n, H.dump_samples_dir, batch_size=batch_size) # ── 1. Sample Grid ── print("\n[1/3] Generating sample grid...") grid_path = os.path.join(output_dir, "samples_grid.png") generate_grid_image(eval_model, H.latent_dim, H.num_grid_samples, H.grid_nrow, grid_path) # ── 2. FID ── if not H.skip_fid: print(f"\n[2/3] Computing FID ({H.num_fid_samples} samples)...") fid_dir = os.path.join(output_dir, "fid") os.makedirs(fid_dir, exist_ok=True) t0 = time.time() generate_images_to_dir(eval_model, H.latent_dim, H.num_fid_samples, fid_dir, batch_size=batch_size) ref_dir = f'{H.data_root}/img' print(f" Reference dir: {ref_dir}") print(f" Generated dir: {fid_dir}") 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) else: print("\n[2/3] FID skipped (--skip_fid True)") # ── 3. Precision / Recall ── if not H.skip_pr: print( f"\n[3/3] Computing Precision/Recall ({H.num_pr_samples} samples)...") pr_dir = os.path.join(output_dir, "prec_rec") os.makedirs(pr_dir, exist_ok=True) t0 = time.time() generate_images_to_dir(eval_model, H.latent_dim, H.num_pr_samples, pr_dir, batch_size=batch_size) try: from helpers.improved_precision_recall import compute_prec_recall ref_dir = f'{H.data_root}/img' precision, recall = compute_prec_recall(ref_dir, pr_dir) results['precision'] = precision results['recall'] = recall print(f" Precision = {precision:.4f}") print(f" Recall = {recall:.4f}") print(f" ({time.time() - t0:.1f}s)") except ImportError: print(" WARNING: helpers.improved_precision_recall not found, skipping P/R") shutil.rmtree(pr_dir, ignore_errors=True) else: print("\n[3/3] Precision/Recall skipped (--skip_pr True)") # ── Save results ── # Get final cycle values (after override) model_unwrapped = eval_model.module if hasattr( eval_model, 'module') else eval_model mapper = model_unwrapped.decoder.mapping_network inner = getattr(mapper, 'trm', None) results['config'] = { 'restore_path': H.restore_path, '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, } results_path = os.path.join(output_dir, "test_metrics.json") with open(results_path, "w") as f: json.dump(results, f, indent=2) print(f"\nResults saved to {results_path}") print("\n" + "=" * 60) print("RESULTS SUMMARY") print("=" * 60) if 'fid' in results: print(f" FID: {results['fid']:.4f}") if 'precision' in results: print(f" Precision: {results['precision']:.4f}") if 'recall' in results: print(f" Recall: {results['recall']:.4f}") if 'ada_fid' in results: print(f" Ada FID: {results['ada_fid']:.5f}") if 'ada_sfid' in results: print(f" Ada sFID: {results['ada_sfid']:.5f}") if 'ada_inception_score' in results: print(f" Ada IS: {results['ada_inception_score']:.5f}") if 'ada_precision' in results: print(f" Ada Prec: {results['ada_precision']:.4f}") if 'ada_recall' in results: print(f" Ada Recall:{results['ada_recall']:.4f}") print(f" Grid: {grid_path}") print("=" * 60) if __name__ == "__main__": main()