| 27 | |
| 28 | |
| 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): |
| 54 | for opt in opt_list: |
| 55 | opt.zero_grad() |
| 56 | |
| 57 | @staticmethod |
| 58 | def clip_norm(network_list): |
| 59 | for network in network_list: |
| 60 | clip_grad_norm_(network.parameters(), 0.5) |
| 61 | |
| 62 | @staticmethod |
| 63 | def step(opt_list): |
| 64 | for opt in opt_list: |
| 65 | opt.step() |
| 66 | |
| 67 | def forward(self, batch_data, eval_mode=False): |
| 68 | caption, motions, m_lens = batch_data |
| 69 | motions = motions.detach().to(self.device).float() |
| 70 | |
| 71 | self.caption = caption |
| 72 | self.motions = motions |
| 73 | x_start = motions |
| 74 | B, T = x_start.shape[:2] |
| 75 | cur_len = torch.LongTensor([min(T, m_len) for m_len in m_lens]).to(self.device) |
| 76 | t, _ = self.sampler.sample(B, x_start.device) |
| 77 | output = self.diffusion.training_losses( |
| 78 | model=self.encoder, |
| 79 | x_start=x_start, |
| 80 | t=t, |
| 81 | model_kwargs={"text": caption, "length": cur_len} |
| 82 | ) |
| 83 | |
| 84 | self.real_noise = output['target'] |
| 85 | self.fake_noise = output['pred'] |
| 86 | try: |
no outgoing calls
no test coverage detected