Method log_local

examples/images/diffusion/main.py:432–454

(
        self,
        save_dir,
        split,  # The split (train/val) on which to log the images
        images,  # A dictionary of images to log
        global_step,  # The global step
        current_epoch,  # The current epoch.
        batch_idx,
    )

    @rank_zero_only
    def log_local(
        self,
        save_dir,
        split,  # The split (train/val) on which to log the images
        images,  # A dictionary of images to log
        global_step,  # The global step
        current_epoch,  # The current epoch.
        batch_idx,
    ):
        # Method for saving image grids to local file system
        root = os.path.join(save_dir, "images", split)
        for k in images:
            grid = torchvision.utils.make_grid(images[k], nrow=4)
            if self.rescale:
                grid = (grid + 1.0) / 2.0  # -1,1 -> 0,1; c,h,w
            grid = grid.transpose(0, 1).transpose(1, 2).squeeze(-1)
            grid = grid.numpy()
            grid = (grid * 255).astype(np.uint8)
            filename = "{}_gs-{:06}_e-{:06}_b-{:06}.png".format(k, global_step, current_epoch, batch_idx)
            path = os.path.join(root, filename)
            os.makedirs(os.path.split(path)[0], exist_ok=True)
            # Save image grid as PNG file
            Image.fromarray(grid).save(path)

    def log_img(self, pl_module, batch, batch_idx, split="train"):
        # Function for logging images to both the logger and local file system.