| 13 | |
| 14 | |
| 15 | class VAEXperiment(pl.LightningModule): |
| 16 | |
| 17 | def __init__(self, |
| 18 | vae_model: BaseVAE, |
| 19 | params: dict) -> None: |
| 20 | super(VAEXperiment, self).__init__() |
| 21 | |
| 22 | self.model = vae_model |
| 23 | self.params = params |
| 24 | self.curr_device = None |
| 25 | self.hold_graph = False |
| 26 | try: |
| 27 | self.hold_graph = self.params['retain_first_backpass'] |
| 28 | except: |
| 29 | pass |
| 30 | |
| 31 | def forward(self, input: Tensor, **kwargs) -> Tensor: |
| 32 | return self.model(input, **kwargs) |
| 33 | |
| 34 | def training_step(self, batch, batch_idx, optimizer_idx = 0): |
| 35 | real_img, labels = batch |
| 36 | self.curr_device = real_img.device |
| 37 | |
| 38 | results = self.forward(real_img, labels = labels) |
| 39 | train_loss = self.model.loss_function(*results, |
| 40 | M_N = self.params['kld_weight'], #al_img.shape[0]/ self.num_train_imgs, |
| 41 | optimizer_idx=optimizer_idx, |
| 42 | batch_idx = batch_idx) |
| 43 | |
| 44 | self.log_dict({key: val.item() for key, val in train_loss.items()}, sync_dist=True) |
| 45 | |
| 46 | return train_loss['loss'] |
| 47 | |
| 48 | def validation_step(self, batch, batch_idx, optimizer_idx = 0): |
| 49 | real_img, labels = batch |
| 50 | self.curr_device = real_img.device |
| 51 | |
| 52 | results = self.forward(real_img, labels = labels) |
| 53 | val_loss = self.model.loss_function(*results, |
| 54 | M_N = 1.0, #real_img.shape[0]/ self.num_val_imgs, |
| 55 | optimizer_idx = optimizer_idx, |
| 56 | batch_idx = batch_idx) |
| 57 | |
| 58 | self.log_dict({f"val_{key}": val.item() for key, val in val_loss.items()}, sync_dist=True) |
| 59 | |
| 60 | |
| 61 | def on_validation_end(self) -> None: |
| 62 | self.sample_images() |
| 63 | |
| 64 | def sample_images(self): |
| 65 | # Get sample reconstruction image |
| 66 | test_input, test_label = next(iter(self.trainer.datamodule.test_dataloader())) |
| 67 | test_input = test_input.to(self.curr_device) |
| 68 | test_label = test_label.to(self.curr_device) |
| 69 | |
| 70 | # test_input, test_label = batch |
| 71 | recons = self.model.generate(test_input, labels = test_label) |
| 72 | vutils.save_image(recons.data, |