(self, args, encoder)
| 29 | class DDPMTrainer(object): |
| 30 | |
| 31 | def __init__(self, args, encoder): |
| 32 | self.opt = args |
| 33 | self.device = args.device |
| 34 | self.encoder = encoder |
| 35 | self.diffusion_steps = args.diffusion_steps |
| 36 | sampler = 'uniform' |
| 37 | beta_scheduler = 'linear' |
| 38 | betas = get_named_beta_schedule(beta_scheduler, self.diffusion_steps) |
| 39 | self.diffusion = GaussianDiffusion( |
| 40 | betas=betas, |
| 41 | model_mean_type=ModelMeanType.EPSILON, |
| 42 | model_var_type=ModelVarType.FIXED_SMALL, |
| 43 | loss_type=LossType.MSE |
| 44 | ) |
| 45 | self.sampler = create_named_schedule_sampler(sampler, self.diffusion) |
| 46 | self.sampler_name = sampler |
| 47 | |
| 48 | if args.is_train: |
| 49 | self.mse_criterion = torch.nn.MSELoss(reduction='none') |
| 50 | self.to(self.device) |
| 51 | |
| 52 | @staticmethod |
| 53 | def zero_grad(opt_list): |
nothing calls this directly
no test coverage detected