| 523 | |
| 524 | @rank_zero_only |
| 525 | def log_local(self, save_dir, split, images, |
| 526 | global_step, current_epoch, batch_idx): |
| 527 | root = os.path.join(save_dir, "images", split) |
| 528 | for k in images: |
| 529 | grid = torchvision.utils.make_grid(images[k], nrow=4) |
| 530 | if self.rescale: |
| 531 | grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w |
| 532 | grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) |
| 533 | grid = grid.numpy() |
| 534 | grid = (grid * 255).astype(np.uint8) |
| 535 | filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format( |
| 536 | k, |
| 537 | global_step, |
| 538 | current_epoch, |
| 539 | batch_idx) |
| 540 | path = os.path.join(root, filename) |
| 541 | os.makedirs(os.path.split(path)[0], exist_ok=True) |
| 542 | Image.fromarray(grid).save(path) |
| 543 | |
| 544 | def log_img(self, pl_module, batch, batch_idx, split="train"): |
| 545 | check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step |