import os import numpy as np import torch import torchvision from PIL import Image from pytorch_lightning.callbacks import Callback import pytorch_lightning as pl from pytorch_lightning.utilities.distributed import rank_zero_only from omegaconf import OmegaConf # class ImageLogger(Callback): # def __init__(self, batch_frequency=2000, max_images=4, clamp=True, increase_log_steps=True, # rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False, # log_images_kwargs=None): # super().__init__() # self.rescale = rescale # self.batch_freq = batch_frequency # self.max_images = max_images # if not increase_log_steps: # self.log_steps = [self.batch_freq] # self.clamp = clamp # self.disabled = disabled # self.log_on_batch_idx = log_on_batch_idx # self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} # self.log_first_step = log_first_step # @rank_zero_only # def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): # root = os.path.join(save_dir, "image_log", 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) # Image.fromarray(grid).save(path) # def log_img(self, pl_module, batch, batch_idx, split="train"): # check_idx = batch_idx # if self.log_on_batch_idx else pl_module.global_step # if (self.check_frequency(check_idx) and # batch_idx % self.batch_freq == 0 # hasattr(pl_module, "log_images") and # callable(pl_module.log_images) and # self.max_images > 0): # logger = type(pl_module.logger) # is_train = pl_module.training # if is_train: # pl_module.eval() # with torch.no_grad(): # images = pl_module.log_images(batch, split=split, **self.log_images_kwargs) # 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., 1.) # self.log_local(pl_module.logger.save_dir, split, images, # pl_module.global_step, pl_module.current_epoch, batch_idx) # if is_train: # pl_module.train() # def check_frequency(self, check_idx): # return check_idx % self.batch_freq == 0 # def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): # if not self.disabled: # self.log_img(pl_module, batch, batch_idx, split="train") class SetupCallback(Callback): def __init__(self, resume, now, logdir, ckptdir, cfgdir, config, lightning_config): super().__init__() self.resume = resume self.now = now self.logdir = logdir self.ckptdir = ckptdir self.cfgdir = cfgdir self.config = config self.lightning_config = lightning_config def on_keyboard_interrupt(self, trainer, pl_module): if trainer.global_rank == 0: print("Summoning checkpoint.") ckpt_path = os.path.join(self.ckptdir, "last.ckpt") trainer.save_checkpoint(ckpt_path) def on_pretrain_routine_start(self, trainer, pl_module): if trainer.global_rank == 0: # Create logdirs and save configs os.makedirs(self.logdir, exist_ok=True) os.makedirs(self.ckptdir, exist_ok=True) os.makedirs(self.cfgdir, exist_ok=True) if "callbacks" in self.lightning_config: if 'metrics_over_trainsteps_checkpoint' in self.lightning_config['callbacks']: os.makedirs(os.path.join(self.ckptdir, 'trainstep_checkpoints'), exist_ok=True) print("Project config") print(OmegaConf.to_yaml(self.config)) OmegaConf.save(self.config, os.path.join(self.cfgdir, "{}-project.yaml".format(self.now))) print("Lightning config") print(OmegaConf.to_yaml(self.lightning_config)) OmegaConf.save(OmegaConf.create({"lightning": self.lightning_config}), os.path.join(self.cfgdir, "{}-lightning.yaml".format(self.now))) else: # ModelCheckpoint callback created log directory --- remove it if not self.resume and os.path.exists(self.logdir): dst, name = os.path.split(self.logdir) dst = os.path.join(dst, "child_runs", name) os.makedirs(os.path.split(dst)[0], exist_ok=True) try: os.rename(self.logdir, dst) except FileNotFoundError: pass class ImageLogger(Callback): def __init__(self, batch_frequency, max_images, clamp=True, increase_log_steps=True, rescale=True, disabled=False, log_on_batch_idx=False, log_first_step=False, log_images_kwargs=None): super().__init__() self.rescale = rescale self.batch_freq = batch_frequency self.max_images = max_images self.logger_log_images = { pl.loggers.TestTubeLogger: self._testtube, } self.log_steps = [2 ** n for n in range(int(np.log2(self.batch_freq)) + 1)] if not increase_log_steps: self.log_steps = [self.batch_freq] self.clamp = clamp self.disabled = disabled self.log_on_batch_idx = log_on_batch_idx self.log_images_kwargs = log_images_kwargs if log_images_kwargs else {} self.log_first_step = log_first_step @rank_zero_only def _testtube(self, pl_module, images, batch_idx, split): for k in images: grid = torchvision.utils.make_grid(images[k]) grid = (grid + 1.0) / 2.0 # -1,1 -> 0,1; c,h,w tag = f"{split}/{k}" pl_module.logger.experiment.add_image( tag, grid, global_step=pl_module.global_step) @rank_zero_only def log_local(self, save_dir, split, images, global_step, current_epoch, batch_idx): 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) Image.fromarray(grid).save(path) def log_img(self, pl_module, batch, batch_idx, split="train"): check_idx = batch_idx if self.log_on_batch_idx else pl_module.global_step if (self.check_frequency(check_idx) and # batch_idx % self.batch_freq == 0 hasattr(pl_module, "log_images") and callable(pl_module.log_images) and self.max_images > 0): logger = type(pl_module.logger) is_train = pl_module.training if is_train: pl_module.eval() with torch.no_grad(): images = pl_module.log_images(batch, split=split, **self.log_images_kwargs) 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., 1.) self.log_local(pl_module.logger.save_dir, split, images, pl_module.global_step, pl_module.current_epoch, batch_idx) logger_log_images = self.logger_log_images.get(logger, lambda *args, **kwargs: None) logger_log_images(pl_module, images, pl_module.global_step, split) if is_train: pl_module.train() def check_frequency(self, check_idx): if ((check_idx % self.batch_freq) == 0 or (check_idx in self.log_steps)) and ( check_idx > 0 or self.log_first_step): try: self.log_steps.pop(0) except IndexError as e: print(e) pass return True return False def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): if not self.disabled and (pl_module.global_step > 0 or self.log_first_step): self.log_img(pl_module, batch, batch_idx, split="train") def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): # if not self.disabled and pl_module.global_step > 0: # self.log_img(pl_module, batch, batch_idx, split="val") # if hasattr(pl_module, 'calibrate_grad_norm'): # if (pl_module.calibrate_grad_norm and batch_idx % 25 == 0) and batch_idx > 0: # self.log_gradients(trainer, pl_module, batch_idx=batch_idx) pass