| """evaluation.py — Evaluation metrics for NSGF/NSGF++ experiments. |
| |
| Implements: |
| - 2-Wasserstein distance (2D experiments) |
| - FID (Fréchet Inception Distance) for image experiments |
| - IS (Inception Score) for image experiments |
| - Visualization utilities |
| |
| Reference: arXiv:2401.14069, Section 5, Appendix E |
| """ |
|
|
| import os |
| import logging |
| import numpy as np |
| import torch |
| import torch.nn as nn |
| from typing import Dict, Optional, List, Tuple |
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| def compute_w2_distance(samples: torch.Tensor, targets: torch.Tensor) -> float: |
| """Compute 2-Wasserstein distance using POT library.""" |
| import ot |
| x = samples.detach().cpu().numpy() |
| y = targets.detach().cpu().numpy() |
| M = ot.dist(x, y, metric="sqeuclidean") |
| a = np.ones(len(x)) / len(x) |
| b = np.ones(len(y)) / len(y) |
| w2_sq = ot.emd2(a, b, M) |
| return float(np.sqrt(max(w2_sq, 0))) |
|
|
|
|
| class InceptionV3Features(nn.Module): |
| """Inception V3 wrapper for FID/IS computation.""" |
| def __init__(self, device: str = "cpu"): |
| super().__init__() |
| import torchvision.models as models |
| self.device = device |
| inception = models.inception_v3(pretrained=True, transform_input=False) |
| inception.eval() |
| self.blocks = nn.Sequential( |
| inception.Conv2d_1a_3x3, inception.Conv2d_2a_3x3, |
| inception.Conv2d_2b_3x3, nn.MaxPool2d(3, stride=2), |
| inception.Conv2d_3b_1x1, inception.Conv2d_4a_3x3, |
| nn.MaxPool2d(3, stride=2), |
| inception.Mixed_5b, inception.Mixed_5c, inception.Mixed_5d, |
| inception.Mixed_6a, inception.Mixed_6b, inception.Mixed_6c, |
| inception.Mixed_6d, inception.Mixed_6e, |
| inception.Mixed_7a, inception.Mixed_7b, inception.Mixed_7c, |
| ) |
| self.avgpool = nn.AdaptiveAvgPool2d((1, 1)) |
| self.fc = inception.fc |
| self.to(device) |
| for p in self.parameters(): |
| p.requires_grad_(False) |
|
|
| @torch.no_grad() |
| def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: |
| if x.shape[2] != 299 or x.shape[3] != 299: |
| x = torch.nn.functional.interpolate(x, size=(299, 299), mode="bilinear", align_corners=False) |
| if x.shape[1] == 1: |
| x = x.repeat(1, 3, 1, 1) |
| x = (x + 1) / 2 |
| h = self.blocks(x) |
| features = self.avgpool(h).squeeze(-1).squeeze(-1) |
| logits = self.fc(features) |
| return features, logits |
|
|
|
|
| def compute_fid(generated: torch.Tensor, real: torch.Tensor, |
| device: str = "cpu", batch_size: int = 64) -> float: |
| from scipy import linalg |
| model = InceptionV3Features(device) |
| def get_features(images): |
| feats = [] |
| for i in range(0, len(images), batch_size): |
| batch = images[i:i + batch_size].to(device) |
| f, _ = model(batch) |
| feats.append(f.cpu().numpy()) |
| return np.concatenate(feats, axis=0) |
| logger.info("Computing FID: extracting generated features...") |
| feats_gen = get_features(generated) |
| logger.info("Computing FID: extracting real features...") |
| feats_real = get_features(real) |
| mu_gen, sigma_gen = feats_gen.mean(0), np.cov(feats_gen, rowvar=False) |
| mu_real, sigma_real = feats_real.mean(0), np.cov(feats_real, rowvar=False) |
| diff = mu_gen - mu_real |
| covmean, _ = linalg.sqrtm(sigma_gen @ sigma_real, disp=False) |
| if np.iscomplexobj(covmean): |
| covmean = covmean.real |
| fid = diff @ diff + np.trace(sigma_gen + sigma_real - 2 * covmean) |
| return float(fid) |
|
|
|
|
| def compute_inception_score(images: torch.Tensor, device: str = "cpu", |
| batch_size: int = 64, splits: int = 10) -> Tuple[float, float]: |
| model = InceptionV3Features(device) |
| all_logits = [] |
| for i in range(0, len(images), batch_size): |
| batch = images[i:i + batch_size].to(device) |
| _, logits = model(batch) |
| all_logits.append(logits.cpu()) |
| all_logits = torch.cat(all_logits, dim=0) |
| probs = torch.softmax(all_logits, dim=1).numpy() |
| scores = [] |
| n = len(probs) |
| split_size = n // splits |
| for i in range(splits): |
| part = probs[i * split_size:(i + 1) * split_size] |
| py = part.mean(axis=0, keepdims=True) |
| kl = part * (np.log(part + 1e-10) - np.log(py + 1e-10)) |
| kl = kl.sum(axis=1).mean() |
| scores.append(np.exp(kl)) |
| return float(np.mean(scores)), float(np.std(scores)) |
|
|
|
|
| class Evaluation: |
| def __init__(self, config: dict, device: str = "cpu"): |
| self.config = config |
| self.device = device |
| self.dataset_name = config.get("dataset", "8gaussians") |
| self.is_image = self.dataset_name in ("mnist", "cifar10") |
|
|
| def evaluate(self, generated: torch.Tensor, real: torch.Tensor) -> Dict[str, float]: |
| metrics = {} |
| if self.is_image: |
| eval_cfg = self.config.get("evaluation", {}) |
| metric_names = eval_cfg.get("metrics", ["fid"]) |
| if "fid" in metric_names: |
| logger.info("Computing FID...") |
| metrics["fid"] = compute_fid(generated, real, self.device) |
| logger.info(f"FID: {metrics['fid']:.2f}") |
| if "is" in metric_names: |
| logger.info("Computing Inception Score...") |
| is_mean, is_std = compute_inception_score(generated, self.device) |
| metrics["is_mean"] = is_mean |
| metrics["is_std"] = is_std |
| logger.info(f"IS: {is_mean:.2f} ± {is_std:.2f}") |
| else: |
| w2 = compute_w2_distance(generated, real) |
| metrics["w2"] = w2 |
| logger.info(f"W2 distance: {w2:.4f}") |
| return metrics |
|
|
|
|
| def plot_2d_samples(samples: torch.Tensor, targets: torch.Tensor, |
| title: str = "Generated vs Target", save_path: Optional[str] = None): |
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| fig, axes = plt.subplots(1, 3, figsize=(15, 5)) |
| s = samples.detach().cpu().numpy() |
| t = targets.detach().cpu().numpy() |
| axes[0].scatter(t[:, 0], t[:, 1], s=3, alpha=0.5, c="blue") |
| axes[0].set_title("Target Distribution") |
| axes[0].set_xlim(-6, 6); axes[0].set_ylim(-6, 6); axes[0].set_aspect("equal") |
| axes[1].scatter(s[:, 0], s[:, 1], s=3, alpha=0.5, c="red") |
| axes[1].set_title("Generated Samples") |
| axes[1].set_xlim(-6, 6); axes[1].set_ylim(-6, 6); axes[1].set_aspect("equal") |
| axes[2].scatter(t[:, 0], t[:, 1], s=3, alpha=0.3, c="blue", label="Target") |
| axes[2].scatter(s[:, 0], s[:, 1], s=3, alpha=0.3, c="red", label="Generated") |
| axes[2].set_title("Overlay") |
| axes[2].set_xlim(-6, 6); axes[2].set_ylim(-6, 6); axes[2].set_aspect("equal") |
| axes[2].legend() |
| plt.suptitle(title) |
| plt.tight_layout() |
| if save_path: |
| os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") |
| logger.info(f"Saved plot to {save_path}") |
| plt.close() |
|
|
|
|
| def plot_2d_trajectory(trajectory: List[torch.Tensor], targets: torch.Tensor, |
| title: str = "Flow Trajectory", save_path: Optional[str] = None, |
| max_particles: int = 200): |
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| from matplotlib.collections import LineCollection |
| fig, ax = plt.subplots(1, 1, figsize=(8, 8)) |
| t = targets.detach().cpu().numpy() |
| ax.scatter(t[:, 0], t[:, 1], s=3, alpha=0.2, c="blue", label="Target") |
| T = len(trajectory) |
| n = min(trajectory[0].shape[0], max_particles) |
| for i in range(n): |
| points = np.array([trajectory[step][i].detach().cpu().numpy() for step in range(T)]) |
| segments = np.array([[points[j], points[j + 1]] for j in range(len(points) - 1)]) |
| colors = plt.cm.coolwarm(np.linspace(0, 1, len(segments))) |
| lc = LineCollection(segments, colors=colors, linewidths=0.5, alpha=0.5) |
| ax.add_collection(lc) |
| final = trajectory[-1][:n].detach().cpu().numpy() |
| ax.scatter(final[:, 0], final[:, 1], s=5, c="red", alpha=0.5, label="Generated") |
| ax.set_xlim(-6, 6); ax.set_ylim(-6, 6); ax.set_aspect("equal") |
| ax.set_title(title); ax.legend() |
| if save_path: |
| os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") |
| logger.info(f"Saved trajectory plot to {save_path}") |
| plt.close() |
|
|
|
|
| def plot_image_grid(images: torch.Tensor, nrow: int = 8, |
| title: str = "Generated Images", save_path: Optional[str] = None): |
| import matplotlib |
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| import torchvision.utils as vutils |
| grid = vutils.make_grid(images[:nrow * nrow], nrow=nrow, normalize=True, value_range=(-1, 1)) |
| grid_np = grid.permute(1, 2, 0).cpu().numpy() |
| fig, ax = plt.subplots(1, 1, figsize=(10, 10)) |
| ax.imshow(grid_np); ax.set_title(title); ax.axis("off") |
| if save_path: |
| os.makedirs(os.path.dirname(save_path) or ".", exist_ok=True) |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") |
| logger.info(f"Saved image grid to {save_path}") |
| plt.close() |
|
|
