File size: 2,624 Bytes
41dd7e0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 | """EDMMethod — wraps the math in core.py."""
import torch
from ..base import Method
from . import core
class EDMMethod(Method):
"""EDM (Karras et al. 2022) for conditional endpoint prediction.
Hyperparameters
---------------
sigma_data : float
EDM σ_data; matches the empirical per-feature std of the (normalized)
training x_T. For data normalized to unit variance per feature this is
1.0.
P_mean, P_std : float
Log-normal training-σ sampler (Table 1 of the EDM paper).
sampling : dict
Keys: n_steps, sigma_min, sigma_max, rho. Used by Heun's sampler.
"""
def __init__(self, sigma_data: float, P_mean: float, P_std: float, sampling: dict,
optimizer_cfg: dict | None = None):
self.sigma_data = float(sigma_data)
self.P_mean = float(P_mean)
self.P_std = float(P_std)
self.n_steps = int(sampling["n_steps"])
self.sigma_min = float(sampling["sigma_min"])
self.sigma_max = float(sampling["sigma_max"])
self.rho = float(sampling["rho"])
# Picked up by base Method.configure_optimizer; defaults to plain Adam.
self.optimizer_cfg = optimizer_cfg or {}
@classmethod
def from_config(cls, cfg: dict) -> "EDMMethod":
return cls(
sigma_data=cfg["sigma_data"],
P_mean=cfg["P_mean"],
P_std=cfg["P_std"],
sampling=cfg["sampling"],
optimizer_cfg=cfg.get("optimizer"),
)
def loss(self, model: torch.nn.Module, batch) -> torch.Tensor:
x_0, x_T = batch
return core.edm_loss(model, x_T, x_0, self.sigma_data, self.P_mean, self.P_std)
def sample(self, model: torch.nn.Module, x_0: torch.Tensor, n_samples: int) -> torch.Tensor:
"""Return Heun-sampled clouds, shape [B, n_samples, D].
x_0 may be either a single conditioning vector of shape (D,) or a
batch of shape (B, D). The output is always (B, n_samples, D).
"""
device = next(model.parameters()).device
if x_0.ndim == 1:
x_0 = x_0.unsqueeze(0)
B = x_0.shape[0]
clouds = []
for i in range(B):
cloud = core.heun_sample(
model,
x_0[i],
n_samples=n_samples,
N=self.n_steps,
sigma_min=self.sigma_min,
sigma_max=self.sigma_max,
rho=self.rho,
sigma_data=self.sigma_data,
device=device,
)
clouds.append(cloud)
return torch.stack(clouds, dim=0)
|