Upload inference.py

inference.py

@@ -0,0 +1,126 @@
+"""inference.py — Sampling / inference for NSGF and NSGF++.
+
+Implements:
+  - NSGF Euler-step inference (standard model)
+  - NSGF++ two-phase inference (NSGF → phase transition → NSF)
+
+Reference: arXiv:2401.14069, Section 4.4, Appendix D
+"""
+
+import torch
+import torch.nn as nn
+from typing import Optional, Tuple, List
+from dataset_loader import DatasetLoader
+
+
+class NSGFSampler:
+    """Sampler using a trained NSGF velocity field model."""
+    def __init__(self, model: nn.Module, data_loader: DatasetLoader,
+                 num_steps: int = 10, device: str = "cpu"):
+        self.model = model.to(device)
+        self.model.eval()
+        self.data_loader = data_loader
+        self.num_steps = num_steps
+        self.device = device
+
+    @torch.no_grad()
+    def sample(self, n: int) -> torch.Tensor:
+        X = self.data_loader.sample_source(n, self.device)
+        dt = 1.0 / self.num_steps
+        for step in range(self.num_steps):
+            t = torch.full((n,), step * dt, device=self.device)
+            v = self.model(X, t)
+            X = X + dt * v
+        return X
+
+    @torch.no_grad()
+    def sample_trajectory(self, n: int) -> List[torch.Tensor]:
+        X = self.data_loader.sample_source(n, self.device)
+        trajectory = [X.clone()]
+        dt = 1.0 / self.num_steps
+        for step in range(self.num_steps):
+            t = torch.full((n,), step * dt, device=self.device)
+            v = self.model(X, t)
+            X = X + dt * v
+            trajectory.append(X.clone())
+        return trajectory
+
+
+class NSGFPlusPlusSampler:
+    """Sampler for the NSGF++ two-phase model.
+    Phase 1 (NSGF): ≤5 Euler steps with Sinkhorn velocity field
+    Phase 2 (NSF): Straight flow velocity field
+    Total NFE = nsgf_steps + nsf_steps
+    """
+    def __init__(self, nsgf_model: nn.Module, nsf_model: nn.Module,
+                 phase_predictor: Optional[nn.Module], data_loader: DatasetLoader,
+                 nsgf_steps: int = 5, nsf_steps: int = 55, device: str = "cpu"):
+        self.nsgf_model = nsgf_model.to(device)
+        self.nsf_model = nsf_model.to(device)
+        self.nsgf_model.eval()
+        self.nsf_model.eval()
+        if phase_predictor is not None:
+            self.phase_predictor = phase_predictor.to(device)
+            self.phase_predictor.eval()
+        else:
+            self.phase_predictor = None
+        self.data_loader = data_loader
+        self.nsgf_steps = nsgf_steps
+        self.nsf_steps = nsf_steps
+        self.device = device
+
+    @torch.no_grad()
+    def sample(self, n: int) -> torch.Tensor:
+        X = self.data_loader.sample_source(n, self.device)
+        dt_nsgf = 1.0 / self.nsgf_steps
+        for step in range(self.nsgf_steps):
+            t = torch.full((n,), step * dt_nsgf, device=self.device)
+            v = self.nsgf_model(X, t)
+            X = X + dt_nsgf * v
+        if self.phase_predictor is not None:
+            t_start = self.phase_predictor(X)
+        else:
+            t_start = torch.zeros(n, device=self.device)
+        dt_nsf = 1.0 / self.nsf_steps
+        for step in range(self.nsf_steps):
+            t_current = t_start + step * dt_nsf * (1.0 - t_start)
+            t_current = t_current.clamp(0, 1)
+            v = self.nsf_model(X, t_current)
+            X = X + dt_nsf * (1.0 - t_start.view(-1, *([1] * (X.dim() - 1)))) * v
+        return X
+
+    @torch.no_grad()
+    def sample_simple(self, n: int) -> torch.Tensor:
+        """Simplified: NSGF then NSF from t=0 to t=1."""
+        X = self.data_loader.sample_source(n, self.device)
+        dt_nsgf = 1.0 / self.nsgf_steps
+        for step in range(self.nsgf_steps):
+            t = torch.full((n,), step * dt_nsgf, device=self.device)
+            v = self.nsgf_model(X, t)
+            X = X + dt_nsgf * v
+        dt_nsf = 1.0 / self.nsf_steps
+        for step in range(self.nsf_steps):
+            t = torch.full((n,), step * dt_nsf, device=self.device)
+            v = self.nsf_model(X, t)
+            X = X + dt_nsf * v
+        return X
+
+    @torch.no_grad()
+    def sample_trajectory(self, n: int) -> Tuple[List[torch.Tensor], int]:
+        trajectory = []
+        X = self.data_loader.sample_source(n, self.device)
+        trajectory.append(X.clone())
+        dt_nsgf = 1.0 / self.nsgf_steps
+        for step in range(self.nsgf_steps):
+            t = torch.full((n,), step * dt_nsgf, device=self.device)
+            v = self.nsgf_model(X, t)
+            X = X + dt_nsgf * v
+            trajectory.append(X.clone())
+        phase_boundary = len(trajectory) - 1
+        dt_nsf = 1.0 / self.nsf_steps
+        for step in range(self.nsf_steps):
+            t = torch.full((n,), step * dt_nsf, device=self.device)
+            v = self.nsf_model(X, t)
+            X = X + dt_nsf * v
+            trajectory.append(X.clone())
+        return trajectory, phase_boundary