| 546 | |
| 547 | @rank_zero_only |
| 548 | def log_local(self, save_dir, split, images, |
| 549 | global_step, current_epoch, batch_idx): |
| 550 | root = os.path.join(save_dir, "images", split) |
| 551 | for k in images: |
| 552 | grid = torchvision.utils.make_grid(images[k], nrow=4) |
| 553 | if self.rescale: |
| 554 | grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w |
| 555 | grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) |
| 556 | grid = grid.numpy() |
| 557 | grid = (grid * 255).astype(np.uint8) |
| 558 | filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format( |
| 559 | k, |
| 560 | global_step, |
| 561 | current_epoch, |
| 562 | batch_idx) |
| 563 | path = os.path.join(root, filename) |
| 564 | os.makedirs(os.path.split(path)[0], exist_ok=True) |
| 565 | Image.fromarray(grid).save(path) |
| 566 | |
| 567 | def log_img(self, pl_module, batch, batch_idx, split="train"): |
| 568 | check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step |