| import abc |
|
|
| import torch |
| import numpy as np |
|
|
| from sgmse.util.registry import Registry |
|
|
|
|
| PredictorRegistry = Registry("Predictor") |
|
|
|
|
| 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 |
|
|
| @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 |
|
|
| def debug_update_fn(self, x, t, *args): |
| raise NotImplementedError(f"Debug update function not implemented for predictor {self}.") |
|
|
|
|
| @PredictorRegistry.register('euler_maruyama') |
| class EulerMaruyamaPredictor(Predictor): |
| def __init__(self, sde, score_fn, probability_flow=False): |
| super().__init__(sde, score_fn, probability_flow=probability_flow) |
|
|
| def update_fn(self, x, y, t, *args): |
| dt = -1. / self.rsde.N |
| z = torch.randn_like(x) |
| f, g = self.rsde.sde(x, y, t, *args) |
| x_mean = x + f * dt |
| x = x_mean + g[:, None, None, None] * np.sqrt(-dt) * z |
| return x, x_mean |
|
|
|
|
| @PredictorRegistry.register('reverse_diffusion') |
| class ReverseDiffusionPredictor(Predictor): |
| def __init__(self, sde, score_fn, probability_flow=False): |
| super().__init__(sde, score_fn, probability_flow=probability_flow) |
|
|
| def update_fn(self, x, y, t, stepsize): |
| f, g = self.rsde.discretize(x, y, t, stepsize) |
| z = torch.randn_like(x) |
| x_mean = x - f |
| x = x_mean + g[:, None, None, None] * z |
| return x, x_mean |
|
|
|
|
| @PredictorRegistry.register('none') |
| class NonePredictor(Predictor): |
| """An empty predictor that does nothing.""" |
|
|
| def __init__(self, *args, **kwargs): |
| pass |
|
|
| def update_fn(self, x, y, t, *args): |
| return x, x |
|
|
