| import abc
|
| import torch
|
|
|
| from sgmse import sdes
|
| from sgmse.util.registry import Registry
|
|
|
|
|
| CorrectorRegistry = Registry("Corrector")
|
|
|
|
|
| class Corrector(abc.ABC):
|
| """The abstract class for a corrector algorithm."""
|
|
|
| def __init__(self, sde, score_fn, snr, n_steps):
|
| super().__init__()
|
| self.rsde = sde.reverse(score_fn)
|
| self.score_fn = score_fn
|
| self.snr = snr
|
| self.n_steps = n_steps
|
|
|
| @abc.abstractmethod
|
| def update_fn(self, x, y, t, *args):
|
| """One update of the corrector.
|
|
|
| Args:
|
| x: A PyTorch tensor representing the current state
|
| t: A PyTorch tensor representing the current time step.
|
| *args: Possibly additional arguments, in particular `y` for OU processes
|
|
|
| Returns:
|
| x: A PyTorch tensor of the next state.
|
| x_mean: A PyTorch tensor. The next state without random noise. Useful for denoising.
|
| """
|
| pass
|
|
|
|
|
| @CorrectorRegistry.register(name='langevin')
|
| class LangevinCorrector(Corrector):
|
| def __init__(self, sde, score_fn, snr, n_steps):
|
| super().__init__(sde, score_fn, snr, n_steps)
|
| self.score_fn = score_fn
|
| self.n_steps = n_steps
|
| self.snr = snr
|
|
|
| def update_fn(self, x, y, t, *args):
|
| target_snr = self.snr
|
| for _ in range(self.n_steps):
|
| grad = self.score_fn(x, y, t, *args)
|
| noise = torch.randn_like(x)
|
| grad_norm = torch.norm(grad.reshape(grad.shape[0], -1), dim=-1).mean()
|
| noise_norm = torch.norm(noise.reshape(noise.shape[0], -1), dim=-1).mean()
|
| step_size = ((target_snr * noise_norm / grad_norm) ** 2 * 2).unsqueeze(0)
|
| x_mean = x + step_size[:, None, None, None] * grad
|
| x = x_mean + noise * torch.sqrt(step_size * 2)[:, None, None, None]
|
|
|
| return x, x_mean
|
|
|
|
|
| @CorrectorRegistry.register(name='ald')
|
| class AnnealedLangevinDynamics(Corrector):
|
| """The original annealed Langevin dynamics predictor in NCSN/NCSNv2."""
|
| def __init__(self, sde, score_fn, snr, n_steps):
|
| super().__init__(sde, score_fn, snr, n_steps)
|
| self.sde = sde
|
| self.score_fn = score_fn
|
| self.snr = snr
|
| self.n_steps = n_steps
|
|
|
| def update_fn(self, x, y, t, *args):
|
| n_steps = self.n_steps
|
| target_snr = self.snr
|
| std = self.sde.marginal_prob(x, y, t, *args)[1]
|
|
|
| for _ in range(n_steps):
|
| grad = self.score_fn(x, y, t, *args)
|
| noise = torch.randn_like(x)
|
| step_size = (target_snr * std) ** 2 * 2
|
| x_mean = x + step_size[:, None, None, None] * grad
|
| x = x_mean + noise * torch.sqrt(step_size * 2)[:, None, None, None]
|
|
|
| return x, x_mean
|
|
|
|
|
| @CorrectorRegistry.register(name='none')
|
| class NoneCorrector(Corrector):
|
| """An empty corrector that does nothing."""
|
|
|
| def __init__(self, *args, **kwargs):
|
| self.snr = 0
|
| self.n_steps = 0
|
| pass
|
|
|
| def update_fn(self, x, t, *args):
|
| return x, x
|
|
|
