| | 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 |
| |
|
