File size: 3,146 Bytes
05b4fca |
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 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 |
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, 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, t, *args):
target_snr = self.snr
for _ in range(self.n_steps):
grad = self.score_fn(x, 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)
if not isinstance(sde, (sdes.OUVESDE,)):
raise NotImplementedError(f"SDE class {sde.__class__.__name__} not yet supported.")
self.sde = sde
self.score_fn = score_fn
self.snr = snr
self.n_steps = n_steps
def update_fn(self, x, t, *args):
n_steps = self.n_steps
target_snr = self.snr
std = self.sde.marginal_prob(x, t, *args)[1]
for _ in range(n_steps):
grad = self.score_fn(x, 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
|