lars
7f62141445
validate_marginals and collect_samples could already vary flow ODE steps for inference (giant predict --steps), but training-time marginal validation and DDIM evaluation were stuck at hardcoded defaults. Add a validate_steps config/CLI option and forward steps to sample_ddim consistently. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
96 lines
3.1 KiB
Python
96 lines
3.1 KiB
Python
import numpy as np
|
|
import torch
|
|
from torch.utils.data import DataLoader
|
|
|
|
from giant.constants import LOCAL_TARGET_NAMES
|
|
from giant.sample import sample_flow, sample_ddpm, sample_ddim
|
|
|
|
|
|
def _kw(steps: int | None) -> dict[str, int]:
|
|
return {} if steps is None else {"steps": steps}
|
|
|
|
|
|
def _histogram_kl(
|
|
p_samples: np.ndarray, q_samples: np.ndarray, bins: int = 50, eps: float = 1e-8
|
|
) -> float:
|
|
"""KL(P || Q) between two 1D samples, estimated via a shared histogram."""
|
|
lo = min(p_samples.min(), q_samples.min())
|
|
hi = max(p_samples.max(), q_samples.max())
|
|
if hi <= lo:
|
|
return 0.0
|
|
edges = np.linspace(lo, hi, bins + 1)
|
|
p_hist, _ = np.histogram(p_samples, bins=edges)
|
|
q_hist, _ = np.histogram(q_samples, bins=edges)
|
|
p = p_hist.astype(np.float64) + eps
|
|
q = q_hist.astype(np.float64) + eps
|
|
p /= p.sum()
|
|
q /= q.sum()
|
|
return float(np.sum(p * np.log(p / q)))
|
|
|
|
|
|
def validate_marginals(
|
|
model: torch.nn.Module,
|
|
val_loader: DataLoader,
|
|
mode: str = "flow",
|
|
schedule=None,
|
|
device: torch.device | None = None,
|
|
n_batches: int | None = None,
|
|
kl_bins: int = 50,
|
|
steps: int | None = None,
|
|
) -> dict[str, np.ndarray]:
|
|
"""Compare per-dimension marginals of generated vs. real steps.
|
|
|
|
Returns {"real": (N,9), "generated": (N,9), "kl_divergence": (9,)} in
|
|
normalised space. `kl_divergence[j]` is KL(real || generated) for
|
|
dimension j, estimated from a shared histogram over both samples.
|
|
|
|
`steps` overrides the number of sampler steps (flow ODE steps or DDIM
|
|
substeps); `None` keeps each sampler's own default. Unused in "ddpm"
|
|
mode, which always runs the full schedule.
|
|
"""
|
|
if device is None:
|
|
device = next(model.parameters()).device
|
|
model.eval()
|
|
|
|
all_real, all_gen = [], []
|
|
for i, (cond_cont, cond_cat, x1) in enumerate(val_loader):
|
|
if n_batches is not None and i >= n_batches:
|
|
break
|
|
cond_cont = cond_cont.to(device)
|
|
cond_cat = cond_cat.to(device)
|
|
|
|
if mode == "flow":
|
|
gen = sample_flow(model, cond_cont, cond_cat, **_kw(steps))
|
|
elif mode == "ddpm":
|
|
gen = sample_ddpm(model, cond_cont, cond_cat, schedule)
|
|
else:
|
|
gen = sample_ddim(model, cond_cont, cond_cat, schedule, **_kw(steps))
|
|
|
|
all_real.append(x1.numpy())
|
|
all_gen.append(gen.cpu().numpy())
|
|
|
|
real = np.concatenate(all_real, axis=0)
|
|
generated = np.concatenate(all_gen, axis=0)
|
|
|
|
kl_divergence = np.array(
|
|
[
|
|
_histogram_kl(real[:, j], generated[:, j], bins=kl_bins)
|
|
for j in range(real.shape[1])
|
|
]
|
|
)
|
|
|
|
header = (
|
|
f"{'Dim':<20} {'real_mean':>10} {'gen_mean':>10} "
|
|
f"{'real_std':>10} {'gen_std':>10} {'KL(real||gen)':>14}"
|
|
)
|
|
print(f"\n{header}")
|
|
print("-" * len(header))
|
|
for j, name in enumerate(LOCAL_TARGET_NAMES):
|
|
r, g = real[:, j], generated[:, j]
|
|
print(
|
|
f"{name:<20} {r.mean():>10.4f} {g.mean():>10.4f} "
|
|
f"{r.std():>10.4f} {g.std():>10.4f} {kl_divergence[j]:>14.4f}"
|
|
)
|
|
|
|
return {"real": real, "generated": generated, "kl_divergence": kl_divergence}
|