File size: 1,892 Bytes
d4c980e |
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 |
import abc
import torch
from geco import sdes
from geco.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='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, t, y, m):
x_mean = 0
n_steps = self.n_steps
target_snr = self.snr
std = self.sde.marginal_prob(x, t, y)[1]
for _ in range(n_steps):
# print(x.shape, y.shape,m.shape)
grad = self.score_fn(x, t, y, m)
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
|