| 271 | |
| 272 | |
| 273 | class Model(nn.Module): |
| 274 | def __init__(self, args, betas, loss_type: str, model_mean_type: str, model_var_type:str): |
| 275 | super(Model, self).__init__() |
| 276 | self.diffusion = GaussianDiffusion(betas, loss_type, model_mean_type, model_var_type) |
| 277 | |
| 278 | self.model = PVCNN2(num_classes=args.nc, embed_dim=args.embed_dim, use_att=args.attention, |
| 279 | dropout=args.dropout, extra_feature_channels=0) |
| 280 | |
| 281 | def prior_kl(self, x0): |
| 282 | return self.diffusion._prior_bpd(x0) |
| 283 | |
| 284 | def all_kl(self, x0, clip_denoised=True): |
| 285 | total_bpd_b, vals_bt, prior_bpd_b, mse_bt = self.diffusion.calc_bpd_loop(self._denoise, x0, clip_denoised) |
| 286 | |
| 287 | return { |
| 288 | 'total_bpd_b': total_bpd_b, |
| 289 | 'terms_bpd': vals_bt, |
| 290 | 'prior_bpd_b': prior_bpd_b, |
| 291 | 'mse_bt':mse_bt |
| 292 | } |
| 293 | |
| 294 | |
| 295 | def _denoise(self, data, t): |
| 296 | B, D,N= data.shape |
| 297 | assert data.dtype == torch.float |
| 298 | assert t.shape == torch.Size([B]) and t.dtype == torch.int64 |
| 299 | |
| 300 | out = self.model(data, t) |
| 301 | |
| 302 | assert out.shape == torch.Size([B, D, N]) |
| 303 | return out |
| 304 | |
| 305 | def get_loss_iter(self, data, noises=None): |
| 306 | B, D, N = data.shape |
| 307 | t = torch.randint(0, self.diffusion.num_timesteps, size=(B,), device=data.device) |
| 308 | |
| 309 | if noises is not None: |
| 310 | noises[t!=0] = torch.randn((t!=0).sum(), *noises.shape[1:]).to(noises) |
| 311 | |
| 312 | losses = self.diffusion.p_losses( |
| 313 | denoise_fn=self._denoise, data_start=data, t=t, noise=noises) |
| 314 | assert losses.shape == t.shape == torch.Size([B]) |
| 315 | return losses |
| 316 | |
| 317 | def gen_samples(self, shape, device, noise_fn=torch.randn, constrain_fn=lambda x, t:x, |
| 318 | clip_denoised=False, max_timestep=None, |
| 319 | keep_running=False): |
| 320 | return self.diffusion.p_sample_loop(self._denoise, shape=shape, device=device, noise_fn=noise_fn, |
| 321 | constrain_fn=constrain_fn, |
| 322 | clip_denoised=clip_denoised, max_timestep=max_timestep, |
| 323 | keep_running=keep_running) |
| 324 | |
| 325 | def reconstruct(self, x0, t, constrain_fn=lambda x, t:x): |
| 326 | |
| 327 | return self.diffusion.reconstruct(x0, t, self._denoise, constrain_fn=constrain_fn) |
| 328 | |
| 329 | def train(self): |
| 330 | self.model.train() |