Spaces:
Sleeping
Sleeping
File size: 3,881 Bytes
a95f6c0 | 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 98 99 100 101 102 103 104 105 106 107 108 | from tqdm import tqdm
import torch
from testing.Sampler import Sampler
class EulerHeunSampler(Sampler):
def __init__(self, model, diff_params, args):
super().__init__(model, diff_params, args)
#stochasticity parameters
self.Schurn=self.args.tester.sampling_params.Schurn
self.Snoise=self.args.tester.sampling_params.Snoise
self.Stmin=self.args.tester.sampling_params.Stmin
self.Stmax=self.args.tester.sampling_params.Stmax
#order of the sampler
self.order=self.args.tester.sampling_params.order
def initialize_x(self, shape, device, schedule):
x = schedule[0]*torch.randn(shape).to(device)
return x
def get_gamma(self, t):
"""
Get the parameter gamma that defines the stochasticity of the sampler
Args
t (Tensor): shape: (N_steps, ) Tensor of timesteps, from which we will compute gamma
"""
N=t.shape[0]
gamma=torch.zeros(t.shape).to(t.device)
#If desired, only apply stochasticity between a certain range of noises Stmin is 0 by default and Stmax is a huge number by default. (Unless these parameters are specified, this does nothing)
indexes=torch.logical_and(t>self.Stmin , t<self.Stmax)
#We use Schurn=5 as the default in our experiments
gamma[indexes]=gamma[indexes]+torch.min(torch.Tensor([self.Schurn/N, 2**(1/2) -1]))
return gamma
def stochastic_timestep(self, x, t, gamma, Snoise=1):
t_hat = t + gamma*t #if gamma_sig[i]==0 this is a deterministic step, make sure it doed not crash
epsilon = torch.randn(x.shape).to(x.device) * Snoise #sample Gaussiannoise, Snoise is 1 by default
x_hat = x + ((t_hat**2 - t**2)**(1/2)) * epsilon #Perturb data
return x_hat, t_hat
def step(self, x_i, t_i, t_iplus1, gamma_i, blind=False):
with torch.no_grad():
x_hat, t_hat = self.stochastic_timestep(x_i, t_i, gamma_i)
x_den = self.get_Tweedie_estimate(x_hat, t_hat)
score = self.Tweedie2score(x_den, x_hat, t_hat)
ode_integrand = self.diff_params._ode_integrand(x_hat, t_hat, score)
dt = t_iplus1 - t_hat
if t_iplus1 !=0 and self.order==2: #second order correction
t_prime = t_iplus1
x_prime = x_hat + dt * ode_integrand
x_den = self.get_Tweedie_estimate(x_prime, t_prime)
score=self.Tweedie2score(x_den, x_prime, t_prime)
ode_integrand_next = self.diff_params._ode_integrand(x_prime, t_prime, score)
ode_integrand_midpoint = .5 * (ode_integrand + ode_integrand_next)
x_iplus1 = x_hat + dt * ode_integrand_midpoint
else:
x_iplus1 = x_hat + dt * ode_integrand
return x_iplus1, x_den
def predict(
self,
shape, #observations (lowpssed signal) Tensor with shape ??
device, #lambda function
blind=False
):
#get the noise schedule
t = self.create_schedule().to(device)
#sample prior
x = self.initialize_x(shape,device, t)
#parameter for langevin stochasticity, if Schurn is 0, gamma will be 0 to, so the sampler will be deterministic
gamma=self.get_gamma(t).to(device)
for i in tqdm(range(0, self.T, 1)):
self.step_counter=i
x, x_den = self.step(x, t[i] , t[i+1], gamma[i], blind)
return x.detach()
def predict_unconditional(
self,
shape, #observations (lowpssed signal) Tensor with shape ??
device
):
self.y=None
self.degradation=None
return self.predict(shape, device)
def predict_conditional(self, *args, **kwargs):
raise NotImplementedError
|