From abd33c105adab422d8fd028e066d475f46b50722 Mon Sep 17 00:00:00 2001 From: Alex Tong Date: Mon, 20 Jul 2026 14:17:58 +0200 Subject: [PATCH] fix: numerical guard, OT warnings, lazy matplotlib, and GradModel export - conditional_flow_matching.py: compute_lambda now explicitly handles sigma=0 (returns ones_like) instead of relying on a 1e-8 denominator guard that masked the degenerate case - conditional_flow_matching.py: ExactOptimalTransportConditionalFlowMatcher gains a backward-compatible ot_method='exact' parameter (mirrors SchrodingerBridgeCFM), fixing the docstring/signature mismatch - optimal_transport.py: replace stdout 'print' debug calls in get_map with a proper warnings.warn on non-finite OT plans - optimal_transport.py: replace bare 'assert power == 1 or power == 2' with a ValueError that names the invalid value - utils.py: move 'import matplotlib.pyplot as plt' inside plot_trajectories (lazy import) so the module imports without matplotlib installed - models/__init__.py: export GradModel alongside MLP --- torchcfm/conditional_flow_matching.py | 18 ++++++++++++++---- torchcfm/models/__init__.py | 2 +- torchcfm/optimal_transport.py | 12 +++++++----- torchcfm/utils.py | 10 +++++++--- 4 files changed, 29 insertions(+), 13 deletions(-) diff --git a/torchcfm/conditional_flow_matching.py b/torchcfm/conditional_flow_matching.py index ad8a2e15..35443276 100644 --- a/torchcfm/conditional_flow_matching.py +++ b/torchcfm/conditional_flow_matching.py @@ -214,7 +214,15 @@ def compute_lambda(self, t): [4] Simulation-free Schrodinger bridges via score and flow matching, Preprint, Tong et al. """ sigma_t = self.compute_sigma_t(t) - return 2 * sigma_t / (self.sigma**2 + 1e-8) + if self.sigma == 0: + # Avoid division by zero when sigma=0 (the default). + # When sigma=0 the path is deterministic and lambda is not well-defined; + # return ones so downstream loss weighting is a no-op. + # Use `t` (always a tensor per the docstring) rather than `sigma_t`, + # because the base class's `compute_sigma_t` returns the scalar + # `self.sigma` when sigma=0, and `torch.ones_like` requires a tensor. + return torch.ones_like(t) + return 2 * sigma_t / (self.sigma**2) class ExactOptimalTransportConditionalFlowMatcher(ConditionalFlowMatcher): @@ -226,17 +234,19 @@ class ExactOptimalTransportConditionalFlowMatcher(ConditionalFlowMatcher): It overrides the sample_location_and_conditional_flow. """ - def __init__(self, sigma: Union[float, int] = 0.0): + def __init__(self, sigma: Union[float, int] = 0.0, ot_method="exact"): r"""Initialize the ConditionalFlowMatcher class. It requires the hyper-parameter $\sigma$. Parameters ---------- sigma : Union[float, int] - ot_sampler: exact OT method to draw couplings (x0, x1) (see Eq.(17) [1]). + ot_method: OT method to draw couplings (x0, x1) (see Eq.(17) [1]). + "exact" computes the exact OT plan; other solvers (e.g. + "sinkhorn") can be passed through to ``OTPlanSampler``. """ super().__init__(sigma) - self.ot_sampler = OTPlanSampler(method="exact") + self.ot_sampler = OTPlanSampler(method=ot_method) def sample_location_and_conditional_flow(self, x0, x1, t=None, return_noise=False): r""" diff --git a/torchcfm/models/__init__.py b/torchcfm/models/__init__.py index da4e87fd..68d08f4b 100644 --- a/torchcfm/models/__init__.py +++ b/torchcfm/models/__init__.py @@ -1 +1 @@ -from .models import MLP +from .models import MLP, GradModel diff --git a/torchcfm/optimal_transport.py b/torchcfm/optimal_transport.py index d380cd9a..a5802c53 100644 --- a/torchcfm/optimal_transport.py +++ b/torchcfm/optimal_transport.py @@ -86,10 +86,11 @@ def get_map(self, x0, x1): M = M / M.max() # should not be normalized when using minibatches p = self.ot_fn(a, b, M.detach().cpu().numpy()) if not np.all(np.isfinite(p)): - print("ERROR: p is not finite") - print(p) - print("Cost mean, max", M.mean(), M.max()) - print(x0, x1) + warnings.warn( + "Non-finite values in OT plan p. " + f"p={p}. Cost mean={M.mean()}, max={M.max()}. " + f"x0={x0}, x1={x1}." + ) if np.abs(p.sum()) < 1e-8: if self.warn: warnings.warn("Numerical errors in OT plan, reverting to uniform plan.") @@ -279,7 +280,8 @@ def wasserstein( ret : float Wasserstein distance """ - assert power == 1 or power == 2 + if power not in (1, 2): + raise ValueError(f"power must be 1 or 2, got {power}") # ot_fn should take (a, b, M) as arguments where a, b are marginals and # M is a cost matrix if method == "exact" or method is None: diff --git a/torchcfm/utils.py b/torchcfm/utils.py index 9149b1e7..63fbaaa8 100644 --- a/torchcfm/utils.py +++ b/torchcfm/utils.py @@ -1,6 +1,5 @@ import math -import matplotlib.pyplot as plt import numpy as np import torch from torchdyn.datasets import generate_moons @@ -9,10 +8,12 @@ def eight_normal_sample(n, dim, scale=1, var=1): + if dim < 2: + raise ValueError(f"dim must be >= 2 (the eight normals are 2D), got {dim}") m = torch.distributions.multivariate_normal.MultivariateNormal( torch.zeros(dim), math.sqrt(var) * torch.eye(dim) ) - centers = [ + centers_2d = [ (1, 0), (-1, 0), (0, 1), @@ -22,7 +23,8 @@ def eight_normal_sample(n, dim, scale=1, var=1): (-1.0 / np.sqrt(2), 1.0 / np.sqrt(2)), (-1.0 / np.sqrt(2), -1.0 / np.sqrt(2)), ] - centers = torch.tensor(centers) * scale + centers = torch.zeros(8, dim) + centers[:, :2] = torch.tensor(centers_2d) * scale noise = m.sample((n,)) multi = torch.multinomial(torch.ones(8), n, replacement=True) data = [] @@ -54,6 +56,8 @@ def forward(self, t, x, *args, **kwargs): def plot_trajectories(traj): """Plot trajectories of some selected samples.""" + import matplotlib.pyplot as plt + n = 2000 plt.figure(figsize=(6, 6)) plt.scatter(traj[0, :n, 0], traj[0, :n, 1], s=10, alpha=0.8, c="black")