| 400 | |
| 401 | |
| 402 | class Model(nn.Module): |
| 403 | def __init__(self, args, betas, loss_type: str, model_mean_type: str, model_var_type:str): |
| 404 | super(Model, self).__init__() |
| 405 | self.diffusion = GaussianDiffusion(betas, loss_type, model_mean_type, model_var_type) |
| 406 | |
| 407 | self.model = Tiger_Transformer_custom(num_classes=args.nc, embed_dim=args.embed_dim, use_att=args.attention, |
| 408 | dropout=args.dropout, extra_feature_channels=0) |
| 409 | |
| 410 | def prior_kl(self, x0): |
| 411 | return self.diffusion._prior_bpd(x0) |
| 412 | |
| 413 | def all_kl(self, x0, clip_denoised=True): |
| 414 | total_bpd_b, vals_bt, prior_bpd_b, mse_bt = self.diffusion.calc_bpd_loop(self._denoise, x0, clip_denoised) |
| 415 | |
| 416 | return { |
| 417 | 'total_bpd_b': total_bpd_b, |
| 418 | 'terms_bpd': vals_bt, |
| 419 | 'prior_bpd_b': prior_bpd_b, |
| 420 | 'mse_bt':mse_bt |
| 421 | } |
| 422 | |
| 423 | |
| 424 | def _denoise(self, data, t): |
| 425 | B, D,N= data.shape |
| 426 | assert data.dtype == torch.float |
| 427 | assert t.shape == torch.Size([B]) and t.dtype == torch.int64 |
| 428 | |
| 429 | out = self.model(data, t) |
| 430 | |
| 431 | assert out.shape == torch.Size([B, D, N]) |
| 432 | return out |
| 433 | |
| 434 | def get_loss_iter(self, data, noises=None): |
| 435 | B, D, N = data.shape |
| 436 | t = torch.randint(0, self.diffusion.num_timesteps, size=(B,), device=data.device) |
| 437 | |
| 438 | if noises is not None: |
| 439 | noises[t!=0] = torch.randn((t!=0).sum(), *noises.shape[1:]).to(noises) |
| 440 | |
| 441 | losses = self.diffusion.p_losses( |
| 442 | denoise_fn=self._denoise, data_start=data, t=t, noise=noises) |
| 443 | assert losses.shape == t.shape == torch.Size([B]) |
| 444 | return losses |
| 445 | |
| 446 | def gen_samples(self, shape, device, noise_fn=torch.randn, |
| 447 | clip_denoised=True, |
| 448 | keep_running=False): |
| 449 | return self.diffusion.p_sample_loop(self._denoise, shape=shape, device=device, noise_fn=noise_fn, |
| 450 | clip_denoised=clip_denoised, |
| 451 | keep_running=keep_running) |
| 452 | |
| 453 | def gen_sample_traj(self, shape, device, freq, noise_fn=torch.randn, |
| 454 | clip_denoised=True,keep_running=False): |
| 455 | return self.diffusion.p_sample_loop_trajectory(self._denoise, shape=shape, device=device, noise_fn=noise_fn, freq=freq, |
| 456 | clip_denoised=clip_denoised, |
| 457 | keep_running=keep_running) |
| 458 | |
| 459 | def train(self): |