(self, shape, data_batch, num_samp,
aux_info={},
verbose=True,
return_diffusion=False,
return_guidance_losses=False,
class_free_guide_w=0.0,
apply_guidance=True,
guide_clean=False)
| 930 | |
| 931 | @torch.no_grad() |
| 932 | def p_sample_loop(self, shape, data_batch, num_samp, |
| 933 | aux_info={}, |
| 934 | verbose=True, |
| 935 | return_diffusion=False, |
| 936 | return_guidance_losses=False, |
| 937 | class_free_guide_w=0.0, |
| 938 | apply_guidance=True, |
| 939 | guide_clean=False): |
| 940 | device = self.betas.device |
| 941 | |
| 942 | batch_size = shape[0] |
| 943 | if self.current_perturbation_guidance.current_guidance is not None and not apply_guidance: |
| 944 | print('DIFFUSER: Note, not using guidance during sampling, only evaluating guidance loss at very end...') |
| 945 | |
| 946 | # sample from base distribution |
| 947 | x = torch.randn(shape, device=device) # (B, N, T, D) |
| 948 | |
| 949 | x = TensorUtils.join_dimensions(x, begin_axis=0, end_axis=2) # B*N, T, D |
| 950 | |
| 951 | aux_info = TensorUtils.repeat_by_expand_at(aux_info, repeats=num_samp, dim=0) |
| 952 | if return_diffusion: diffusion = [x] |
| 953 | progress = Progress(self.n_timesteps) if verbose else Silent() |
| 954 | |
| 955 | steps = [i for i in reversed(range(0, self.n_timesteps, self.stride))] |
| 956 | # print('steps', steps) |
| 957 | for i in steps: |
| 958 | # print('i', i) |
| 959 | timesteps = torch.full((batch_size*num_samp,), i, device=device, dtype=torch.long) |
| 960 | |
| 961 | x, guide_losses = self.p_sample(x, timesteps, data_batch, aux_info=aux_info, num_samp=num_samp, class_free_guide_w=class_free_guide_w, |
| 962 | apply_guidance=apply_guidance, guide_clean=guide_clean, eval_final_guide_loss=(i == steps[-1])) |
| 963 | # apply hard constraints (overwrite waypoints at certain timesteps) |
| 964 | if self.current_constraints is not None: # and i != steps[-1]: # TODO don't do it for last step? |
| 965 | # TODO why isn't this working very well? And why is y upside down? |
| 966 | # apply constraints expects traj in shape (B, N, T, D) and metric space |
| 967 | x = self.descale_traj(x.reshape((shape[0], shape[1], shape[2], -1))) |
| 968 | x = apply_constraints(x, data_batch['scene_index'], self.current_constraints) |
| 969 | x = self.scale_traj(x.reshape((shape[0]*shape[1], shape[2], -1))) |
| 970 | |
| 971 | |
| 972 | progress.update({'t': i}) |
| 973 | |
| 974 | if return_diffusion: diffusion.append(x) |
| 975 | |
| 976 | progress.close() |
| 977 | |
| 978 | if any(guide_losses): |
| 979 | print('===== GUIDANCE LOSSES ======') |
| 980 | for k,v in guide_losses.items(): |
| 981 | print('%s: %.012f' % (k, np.nanmean(v.cpu()))) |
| 982 | |
| 983 | x = TensorUtils.reshape_dimensions(x, begin_axis=0, end_axis=1, target_dims=(batch_size, num_samp)) |
| 984 | |
| 985 | out_dict = {'pred_traj' : x} |
| 986 | if return_guidance_losses: |
| 987 | out_dict['guide_losses'] = guide_losses |
| 988 | if return_diffusion: |
| 989 | diffusion = [TensorUtils.reshape_dimensions(cur_diff, begin_axis=0, end_axis=1, target_dims=(batch_size, num_samp)) |
no test coverage detected