Spaces:
Sleeping
Sleeping
| 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 | |