| from __future__ import annotations |
| import torch |
| from ..base import Method |
| from . import core |
|
|
| class CascadedSIMethod(Method): |
| def __init__(self, sigma_train, t_eps, sampling, beta_power=1, |
| eval_noise_scales=None, optimizer_cfg=None): |
| self.sigma_train = float(sigma_train) |
| self.t_eps = float(t_eps) |
| self.beta_power = int(beta_power) |
| self.n_sde_steps_per_seg = int(sampling["n_sde_steps_per_seg"]) |
| self.n_segments = int(sampling.get("n_segments", 10)) |
| self.noise_scale = float(sampling.get("noise_scale", 1.0)) |
| |
| |
| if "t_values" in sampling: |
| self.t_values = list(sampling["t_values"]) |
| else: |
| self.t_values = [k / self.n_segments for k in range(self.n_segments + 1)] |
| self.eval_noise_scales = ([float(s) for s in eval_noise_scales] |
| if eval_noise_scales is not None else [self.noise_scale]) |
| self.optimizer_cfg = optimizer_cfg or {} |
|
|
| @classmethod |
| def from_config(cls, cfg): |
| if cfg.get("anchor_ic", False): |
| raise ValueError("This release supports the paper's SI models without retained initial-state inputs") |
| return cls( |
| sigma_train=cfg["sigma_train"], |
| t_eps=cfg.get("t_eps", 0.0), |
| sampling=cfg["sampling"], |
| beta_power=cfg.get("beta_power", 1), |
| eval_noise_scales=cfg.get("eval_noise_scales"), |
| optimizer_cfg=cfg.get("optimizer"), |
| ) |
|
|
| def variants_for_eval(self): |
| return {f"noise_scale_{_fmt(s)}": {"noise_scale": s} for s in self.eval_noise_scales} |
|
|
| def loss(self, model, batch): |
| x_tk, x_tk1, t_k, t_k1 = batch |
| return core.local_si_loss(model, x_tk, x_tk1, t_k, t_k1, |
| sigma=self.sigma_train, t_eps=self.t_eps, |
| beta_power=self.beta_power) |
|
|
| def sample(self, model, x_0, n_samples, noise_scale=None): |
| if noise_scale is None: |
| noise_scale = self.noise_scale |
| if x_0.ndim == 1: |
| x_0 = x_0.unsqueeze(0) |
| B = x_0.shape[0] |
| out = [] |
| for i in range(B): |
| cloud = core.rollout(model, x_0[i], n_samples=n_samples, |
| t_values=self.t_values, |
| n_sde_steps_per_seg=self.n_sde_steps_per_seg, |
| sigma=self.sigma_train, |
| noise_scale=noise_scale) |
| out.append(cloud) |
| return torch.stack(out, dim=0) |
|
|
| def _fmt(s): |
| if s == int(s): |
| return f"{int(s)}p0" |
| return f"{s}".replace(".", "p") |
|
|
