| """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 |
|
|
|
|
| |
| |
| |
| 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(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) |
| |
| 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) |
| |
| 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') |
|
|
| |
| has_module_prefix = any(k.startswith('module.') |
| for k in state_dict.keys()) |
| if has_module_prefix: |
| |
| 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!") |
|
|
| |
| 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 |
|
|
|
|
| |
| |
| |
|
|
| def main(): |
| |
| 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) |
|
|
| |
| 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)') |
| |
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| 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) |
|
|
| |
| imle_model, ema_model = build_and_load_model(H) |
|
|
| |
| 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_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) |
|
|
| |
| 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) |
|
|
| |
| 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)") |
|
|
| |
| 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)") |
|
|
| |
| |
| 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() |
|
|
