| 321 | |
| 322 | @rank_zero_only |
| 323 | def log_local(self, save_dir, split, images, |
| 324 | global_step, current_epoch, batch_idx): |
| 325 | root = os.path.join(save_dir, "images", split) |
| 326 | for k in images: |
| 327 | grid = torchvision.utils.make_grid(images[k], nrow=4) |
| 328 | if self.rescale: |
| 329 | grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w |
| 330 | grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1) |
| 331 | grid = grid.numpy() |
| 332 | grid = (grid * 255).astype(np.uint8) |
| 333 | filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format( |
| 334 | k, |
| 335 | global_step, |
| 336 | current_epoch, |
| 337 | batch_idx) |
| 338 | path = os.path.join(root, filename) |
| 339 | os.makedirs(os.path.split(path)[0], exist_ok=True) |
| 340 | Image.fromarray(grid).save(path) |
| 341 | |
| 342 | def log_img(self, pl_module, batch, batch_idx, split="train"): |
| 343 | check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step |