Method log_img

examples/images/diffusion/main.py:456–497

(self, pl_module, batch, batch_idx, split="train")

            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.
        check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step
        # check if it's time to log an image batch
        if (
            self.check_frequency(check_idx)
            and hasattr(pl_module, "log_images")  # batch_idx % self.batch_freq == 0
            and callable(pl_module.log_images)
            and self.max_images > 0
        ):
            # Get logger type and check if training mode is on
            logger = type(pl_module.logger)

            is_train = pl_module.training
            if is_train:
                pl_module.eval()

            with torch.no_grad():
                # Get images from log_images method of the pl_module
                images = pl_module.log_images(batch, split=split, **self.log_images_kwargs)

            # Clip images if specified and convert to CPU tensor
            for k in images:
                N = min(images[k].shape[0], self.max_images)
                images[k] = images[k][:N]
                if isinstance(images[k], torch.Tensor):
                    images[k] = images[k].detach().cpu()
                    if self.clamp:
                        images[k] = torch.clamp(images[k], -1.0, 1.0)

            # Log images locally to file system
            self.log_local(
                pl_module.logger.save_dir, split, images, pl_module.global_step, pl_module.current_epoch, batch_idx
            )

            # log the images using the logger
            logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None)
            logger_log_images(pl_module, images, pl_module.global_step, split)

            # switch back to training mode if necessary
            if is_train:
                pl_module.train()

    # The function checks if it's time to log an image batch
    def check_frequency(self, check_idx):