(self, pl_module, batch, batch_idx, split="train")
| 51 | blended_image.save(blended_path) |
| 52 | |
| 53 | def log_img(self, pl_module, batch, batch_idx, split="train"): |
| 54 | check_idx = batch_idx # if self.log_on_batch_idx else pl_module.global_step |
| 55 | if (self.check_frequency(check_idx) and # batch_idx % self.batch_freq == 0 |
| 56 | hasattr(pl_module, "log_images") and |
| 57 | callable(pl_module.log_images) and |
| 58 | self.max_images > 0): |
| 59 | logger = type(pl_module.logger) |
| 60 | |
| 61 | is_train = pl_module.training |
| 62 | if is_train: |
| 63 | pl_module.eval() |
| 64 | |
| 65 | with torch.no_grad(): |
| 66 | images = pl_module.log_images(batch, split=split, **self.log_images_kwargs) |
| 67 | |
| 68 | path_image = {} |
| 69 | for k in images: |
| 70 | N = min(images[k].shape[0], self.max_images) |
| 71 | images[k] = images[k][:N] |
| 72 | if isinstance(images[k], torch.Tensor): |
| 73 | images[k] = images[k].detach().cpu() |
| 74 | if self.clamp: |
| 75 | images[k] = torch.clamp(images[k], -1., 1.) |
| 76 | |
| 77 | self.log_local(pl_module.logger.save_dir, split, images, |
| 78 | pl_module.global_step, pl_module.current_epoch, batch_idx) |
| 79 | |
| 80 | if is_train: |
| 81 | pl_module.train() |
| 82 | |
| 83 | def check_frequency(self, check_idx): |
| 84 | return check_idx % self.batch_freq == 0 |
no test coverage detected