(self, network, denoiser, conditioner, input, batch)
| 118 | assert not (self.detach_dcl_prev and self.detach_dcl_post), "Only one detach option can be enabled." |
| 119 | |
| 120 | def __call__(self, network, denoiser, conditioner, input, batch): |
| 121 | cond = conditioner(batch) |
| 122 | additional_model_inputs = {key: batch[key] for key in self.batch2model_keys.intersection(batch)} |
| 123 | |
| 124 | alphas_cumprod_sqrt, idx = self.sigma_sampler(input.shape[0], return_idx=True) |
| 125 | alphas_cumprod_sqrt = alphas_cumprod_sqrt.to(input.device) # a float |
| 126 | idx = idx.to(input.device) # indicating t. shape: [bs]. eg: [845, 10] |
| 127 | |
| 128 | # idx:tensor([26], device='cuda:0'), alpha_t:tensor([0.9631], device='cuda:0') |
| 129 | # idx:tensor([335], device='cuda:0'), alpha_t:tensor([0.5044], device='cuda:0') |
| 130 | # idx:tensor([510], device='cuda:1'), alpha_t:tensor([0.2985], device='cuda:1') |
| 131 | |
| 132 | cond_inds = self.cond_inds |
| 133 | if cond_inds is not None: |
| 134 | |
| 135 | cond_inds = random.choices(cond_inds, weights=self.cond_inds_prob, k=1)[0] # randomly choose a conditioning scheme |
| 136 | cond_inds = list(cond_inds) # convert from omegaconf to normal list |
| 137 | cond_mask = torch.zeros(input.shape).to(input.device) # [1, 13, 16, 64, 112] |
| 138 | cond_mask[:, cond_inds] = 1 |
| 139 | |
| 140 | # TODO: cond_inds is the same for the batch, which might be inoptimal. |
| 141 | # cond_mask[:,:,0,0,0] |
| 142 | # tensor([[1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], |
| 143 | # [1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.]], device='cuda:0') |
| 144 | |
| 145 | noise = torch.randn_like(input) |
| 146 | |
| 147 | # broadcast noise |
| 148 | mp_size = mpu.get_model_parallel_world_size() |
| 149 | global_rank = torch.distributed.get_rank() // mp_size |
| 150 | src = global_rank * mp_size |
| 151 | torch.distributed.broadcast(idx, src=src, group=mpu.get_model_parallel_group()) |
| 152 | torch.distributed.broadcast(noise, src=src, group=mpu.get_model_parallel_group()) |
| 153 | torch.distributed.broadcast(alphas_cumprod_sqrt, src=src, group=mpu.get_model_parallel_group()) |
| 154 | |
| 155 | additional_model_inputs["idx"] = idx |
| 156 | |
| 157 | # Add traj here. |
| 158 | additional_model_inputs["with_traj"] = batch["with_traj"] |
| 159 | additional_model_inputs["fut_traj"] = batch["fut_traj"] |
| 160 | |
| 161 | if self.offset_noise_level > 0.0: |
| 162 | # * Original implementation, but is incorrect |
| 163 | # noise + append_dims(torch.randn(input.shape[0]).to(input.device), input.ndim) * self.offset_noise_level |
| 164 | # ) |
| 165 | b, c, t, _, _ = input.shape |
| 166 | offset_noise_shape = (b, c, t) |
| 167 | |
| 168 | if not self.linear_offset_noise: |
| 169 | # Constant noise level across T |
| 170 | noise = ( |
| 171 | noise + append_dims(torch.randn(offset_noise_shape).to(input.device), input.ndim) * self.offset_noise_level |
| 172 | ) |
| 173 | else: |
| 174 | # linearly increasing noise level across T |
| 175 | offset_noise_level = torch.linspace(0.2, 1, t).to(input.device) * self.offset_noise_level |
| 176 | offset_noise_level = offset_noise_level.unsqueeze(0).unsqueeze(0) # [1, 1, T] |
| 177 | noise = ( |
nothing calls this directly
no test coverage detected