Upload evaluation.py

evaluation.py

@@ -0,0 +1,218 @@
+"""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()