# Adapted from https://github.com/yang-song/score_sde_pytorch/
# and https://github.com/sp-uhh/sgmse
import abc
import numpy as np
import torch
[docs]class Predictor(abc.ABC):
"""The abstract class for a predictor algorithm."""
def __init__(self, sde, score_fn, probability_flow=False):
super().__init__()
self.sde = sde
self.rsde = sde.reverse(score_fn)
self.score_fn = score_fn
self.probability_flow = probability_flow
[docs] @abc.abstractmethod
def update_fn(self, x, t, *args):
"""One update of the predictor.
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
[docs] def debug_update_fn(self, x, t, *args):
raise NotImplementedError(
f"Debug update function not implemented for predictor {self}."
)
[docs]class EulerMaruyamaPredictor(Predictor):
def __init__(self, sde, score_fn, probability_flow=False):
super().__init__(sde, score_fn, probability_flow=probability_flow)
[docs] def update_fn(self, x, t, *args):
dt = -1.0 / self.rsde.N
z = torch.randn_like(x)
f, g = self.rsde.sde(x, t, *args)
x_mean = x + f * dt
x = x_mean + g[:, None, None, None] * np.sqrt(-dt) * z
return x, x_mean
[docs]class ReverseDiffusionPredictor(Predictor):
def __init__(self, sde, score_fn, probability_flow=False):
super().__init__(sde, score_fn, probability_flow=probability_flow)
[docs] def update_fn(self, x, t, *args):
f, g = self.rsde.discretize(x, t, *args)
z = torch.randn_like(x)
x_mean = x - f
x = x_mean + g[:, None, None, None] * z
return x, x_mean
[docs]class NonePredictor(Predictor):
"""An empty predictor that does nothing."""
def __init__(self, *args, **kwargs):
pass
[docs] def update_fn(self, x, t, *args):
return x, x
predictor_dict = dict(
euler_maruyama=EulerMaruyamaPredictor,
reverse_diffusion=ReverseDiffusionPredictor,
none=NonePredictor,
)