# ------------------------------------------------------------------------------------ # Modified from Taming Transformers (https://github.com/CompVis/taming-transformers) # Copyright (c) 2020 Patrick Esser and Robin Rombach and Björn Ommer. All Rights Reserved. # ------------------------------------------------------------------------------------ import os import time import wandb import numpy as np from PIL import Image from pathlib import Path from omegaconf import OmegaConf, DictConfig from typing import Tuple, Generic, Dict, Callable, Optional, Any from pprint import pprint import torch import torchvision import pytorch_lightning as pl import pytorch_lightning.loggers from pytorch_lightning.loggers import WandbLogger from pytorch_lightning.loggers.logger import DummyLogger from pytorch_lightning.utilities import rank_zero_only, rank_zero_info from pytorch_lightning.callbacks import Callback from functools import wraps def node_zero_only(fn: Callable) -> Callable: @wraps(fn) def wrapped_fn(*args, **kwargs) -> Optional[Any]: if node_zero_only.node == 0: return fn(*args, **kwargs) return None return wrapped_fn node_zero_only.node = getattr(node_zero_only, 'node', int(os.environ.get('NODE_RANK', 0))) def node_zero_experiment(fn: Callable) -> Callable: """Returns the real experiment on rank 0 and otherwise the DummyExperiment.""" @wraps(fn) def experiment(self): @node_zero_only def get_experiment(): return fn(self) return get_experiment() or DummyLogger.experiment return experiment # customize wandb for node 0 only class MyWandbLogger(WandbLogger): @WandbLogger.experiment.getter @node_zero_experiment def experiment(self): return super().experiment class SetupCallback(Callback): def __init__(self, config: DictConfig, exp_config: DictConfig, basedir: Path, logdir: str = "log", ckptdir: str = "ckpt") -> None: super().__init__() self.logdir = basedir / logdir self.ckptdir = basedir / ckptdir self.config = config self.exp_config = exp_config # def on_pretrain_routine_start(self, trainer: pl.trainer.Trainer, pl_module: pl.LightningModule) -> None: # if trainer.global_rank == 0: # # Create logdirs and save configs # os.makedirs(self.logdir, exist_ok=True) # os.makedirs(self.ckptdir, exist_ok=True) # # print("Experiment config") # print(self.exp_config.pretty()) # # print("Model config") # print(self.config.pretty()) def on_fit_start(self, trainer: pl.trainer.Trainer, pl_module: pl.LightningModule) -> None: if trainer.global_rank == 0: # Create logdirs and save configs os.makedirs(self.logdir, exist_ok=True) os.makedirs(self.ckptdir, exist_ok=True) # print("Experiment config") # pprint(self.exp_config) # # print("Model config") # pprint(self.config) class ImageLogger(Callback): def __init__(self, batch_frequency: int, max_images: int, clamp: bool = True, increase_log_steps: bool = True) -> None: super().__init__() self.batch_freq = batch_frequency self.max_images = max_images self.logger_log_images = { pl.loggers.WandbLogger: self._wandb, 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 @rank_zero_only def _wandb(self, pl_module, images, batch_idx, split): # raise ValueError("No way wandb") grids = dict() for k in images: grid = torchvision.utils.make_grid(images[k]) grids[f"{split}/{k}"] = wandb.Image(grid) pl_module.logger.experiment.log(grids) @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: str, split: str, images: Dict, global_step: int, current_epoch: int, batch_idx: int) -> None: root = os.path.join(save_dir, "results", split) os.makedirs(root, exist_ok=True) for k in images: grid = torchvision.utils.make_grid(images[k], nrow=4) 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: pl.LightningModule, batch: Tuple[torch.LongTensor, torch.FloatTensor], batch_idx: int, split: str = "train") -> None: if (self.check_frequency(batch_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, pl_module=pl_module) for k in images: N = min(images[k].shape[0], self.max_images) images[k] = images[k][:N].detach().cpu() if self.clamp: images[k] = images[k].clamp(0, 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, batch_idx: int) -> bool: if (batch_idx % self.batch_freq) == 0 or (batch_idx in self.log_steps): try: self.log_steps.pop(0) except IndexError: pass return True return False def on_train_batch_end(self, trainer: pl.trainer.Trainer, pl_module: pl.LightningModule, outputs: Generic, batch: Tuple[torch.LongTensor, torch.FloatTensor], batch_idx: int) -> None: self.log_img(pl_module, batch, batch_idx, split="train") def on_validation_batch_end(self, trainer: pl.trainer.Trainer, pl_module: pl.LightningModule, outputs: Generic, batch: Tuple[torch.LongTensor, torch.FloatTensor], dataloader_idx: int, batch_idx: int) -> None: self.log_img(pl_module, batch, batch_idx, split="val") class CUDACallback(Callback): # see https://github.com/SeanNaren/minGPT/blob/master/mingpt/callback.py def on_train_epoch_start(self, trainer, pl_module): # Reset the memory use counter torch.cuda.reset_peak_memory_stats(trainer.root_gpu) torch.cuda.synchronize(trainer.root_gpu) self.start_time = time.time() def on_train_epoch_end(self, trainer, pl_module, outputs): torch.cuda.synchronize(trainer.root_gpu) max_memory = torch.cuda.max_memory_allocated(trainer.root_gpu) / 2 ** 20 epoch_time = time.time() - self.start_time try: max_memory = trainer.training_type_plugin.reduce(max_memory) epoch_time = trainer.training_type_plugin.reduce(epoch_time) rank_zero_info(f"Average Epoch time: {epoch_time:.2f} seconds") rank_zero_info(f"Average Peak memory {max_memory:.2f}MiB") except AttributeError: pass