File size: 1,754 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 |
import abc
import torch
import numpy as np
from geco.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('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, t, y, m, stepsize):
f, g = self.rsde.discretize(x, t, y, m, stepsize)
z = torch.randn_like(x)
x_mean = x - f
x = x_mean + g[:, None, None, None] * z
return x, x_mean
def update_fn_analyze(self, x, t, *args):
raise NotImplementedError("update_fn_analyze() has not been implemented yet for the ReverseDiffusionPredictor")
|