| """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"]) |
| |
| 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) |
|
|
