(self)
| 172 | self._sample_and_visualize() |
| 173 | |
| 174 | def _sample_and_visualize(self): |
| 175 | print("Sampling and visualizing...") |
| 176 | self.ddp_model.eval() |
| 177 | |
| 178 | batch, cond = next(self.data) |
| 179 | _shape = [2] + list(batch.shape[1:]) |
| 180 | with th.no_grad(): |
| 181 | sample = self.diffusion.p_sample_loop(self.ddp_model, _shape, progress=True, model_kwargs=cond, clip_denoised=False) |
| 182 | sample = sample.detach().cpu().numpy() |
| 183 | feat_dim = sample.shape[1] |
| 184 | for i in range(sample.shape[0]): |
| 185 | fig = draw_scalar_field2D(sample[i, 0]) |
| 186 | self.tb.add_figure(f"sample{i}_c0", fig, global_step=self.step) |
| 187 | fig = draw_scalar_field2D(sample[i, feat_dim // 2]) |
| 188 | self.tb.add_figure(f"sample{i}_c{feat_dim // 2}", fig, global_step=self.step) |
| 189 | |
| 190 | # if self.step == 0: |
| 191 | fig = draw_scalar_field2D(batch[0, 0].detach().cpu().numpy()) |
| 192 | self.tb.add_figure("data_c0", fig, global_step=self.step) |
| 193 | fig = draw_scalar_field2D(batch[0, feat_dim // 2].detach().cpu().numpy()) |
| 194 | self.tb.add_figure(f"data_c{feat_dim // 2}", fig, global_step=self.step) |
| 195 | |
| 196 | self.ddp_model.train() |
| 197 | |
| 198 | def forward_backward(self, batch, cond): |
| 199 | self.mp_trainer.zero_grad() |
no test coverage detected