| 346 | return kl_prior.mean(dim=list(range(1, len(kl_prior.shape)))) / np.log(2.) |
| 347 | |
| 348 | def calc_bpd_loop(self, denoise_fn, x_start, clip_denoised=True): |
| 349 | |
| 350 | with torch.no_grad(): |
| 351 | B, T = x_start.shape[0], self.num_timesteps |
| 352 | |
| 353 | vals_bt_, mse_bt_= torch.zeros([B, T], device=x_start.device), torch.zeros([B, T], device=x_start.device) |
| 354 | for t in reversed(range(T)): |
| 355 | |
| 356 | t_b = torch.empty(B, dtype=torch.int64, device=x_start.device).fill_(t) |
| 357 | # Calculate VLB term at the current timestep |
| 358 | new_vals_b, pred_xstart = self._vb_terms_bpd( |
| 359 | denoise_fn, data_start=x_start, data_t=self.q_sample(x_start=x_start, t=t_b), t=t_b, |
| 360 | clip_denoised=clip_denoised, return_pred_xstart=True) |
| 361 | # MSE for progressive prediction loss |
| 362 | assert pred_xstart.shape == x_start.shape |
| 363 | new_mse_b = ((pred_xstart-x_start)**2).mean(dim=list(range(1, len(x_start.shape)))) |
| 364 | assert new_vals_b.shape == new_mse_b.shape == torch.Size([B]) |
| 365 | # Insert the calculated term into the tensor of all terms |
| 366 | mask_bt = t_b[:, None]==torch.arange(T, device=t_b.device)[None, :].float() |
| 367 | vals_bt_ = vals_bt_ * (~mask_bt) + new_vals_b[:, None] * mask_bt |
| 368 | mse_bt_ = mse_bt_ * (~mask_bt) + new_mse_b[:, None] * mask_bt |
| 369 | assert mask_bt.shape == vals_bt_.shape == vals_bt_.shape == torch.Size([B, T]) |
| 370 | |
| 371 | prior_bpd_b = self._prior_bpd(x_start) |
| 372 | total_bpd_b = vals_bt_.sum(dim=1) + prior_bpd_b |
| 373 | assert vals_bt_.shape == mse_bt_.shape == torch.Size([B, T]) and \ |
| 374 | total_bpd_b.shape == prior_bpd_b.shape == torch.Size([B]) |
| 375 | return total_bpd_b.mean(), vals_bt_.mean(), prior_bpd_b.mean(), mse_bt_.mean() |
| 376 | |
| 377 | |
| 378 | class Tiger_Transformer_custom(Tiger_Transformer): |