File size: 2,071 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 | from __future__ import annotations
import torch
from ..base import Method
from . import core
class DirectSIMethod(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 = int(sampling["n_sde_steps"])
self.noise_scale = float(sampling.get("noise_scale", 1.0))
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):
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_0, x_T = batch
return core.global_si_loss(model, x_0, x_T, 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.sde_sample(model, x_0[i], n_samples=n_samples,
n_steps=self.n_sde_steps, 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")
|