from numpy import sqrt import torch import pytorch_lightning as pl import wandb from argparse import ArgumentParser from kornia.contrib import extract_tensor_patches from torch.optim import Adam, SGD from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts from torchvision.utils import make_grid from src.datasources import S2_ALL_12BANDS from src.losses import mse_loss, mae_loss, ssim_loss, tv_loss, Shift from src.modules import BicubicUpscaledBaseline import pandas as pd def detach_dict(dict): """Detaches the tensors in a dictionary. Parameters ---------- dict : dict of Tensor Dictionary containing the tensors to detach. Returns ------- dict of Tensor Dictionary containing the detached tensors. """ for tensor_key in dict: if dict[tensor_key].requires_grad: dict[tensor_key] = dict[tensor_key].detach() return dict class LitModel(pl.LightningModule): """Lightning Module wrapper. Implements the forward, loss, train/val/test step functions and logging. Parses and adds default model arguments. """ def __init__(self, backbone, **kws) -> None: """Initialises the LightningModule with the backbone and the default arguments. The hyperparameters are fetchd from the hparams attribute. Parameters ---------- backbone : nn.Module Backbone model, as generated by train.py. """ super().__init__() self.save_hyperparameters() # save hyperparams in self.hparams self.backbone = backbone self.transform = self.hparams.transform # transform as a torch module output_channels = [channel for channel in range(self.backbone.out_channels)] self.baseline = BicubicUpscaledBaseline( input_size=self.hparams.input_size, output_size=self.backbone.output_size, output_channels=output_channels, interpolation="bicubic", chip_size=self.hparams.chip_size, ) self.check_loss_functions_weights_are_valid() self.create_shifters() self.create_metrics() self.enable_benchmark_logging() def enable_benchmark_logging(self): if self.hparams.benchmark: self.logging = {} self.logging['current_idx'] = 1 self.logging['log'] = pd.DataFrame() self.log_path = '../../log' def check_loss_functions_weights_are_valid(self): """Checks that the loss functions weights all sum to 1.""" assert self.hparams.w_mse + self.hparams.w_mae + self.hparams.w_ssim == 1 def create_shifters(self): """Creates a shifter for the training and validation steps, and a shifter for the baseline.""" self.shifter = Shift( shift_by_px=self.hparams.shift_px, mode=self.hparams.shift_mode, step=self.hparams.shift_step, use_cache=False, ) self.baseline_shifter = Shift( shift_by_px=self.hparams.shift_px, mode=self.hparams.shift_mode, step=self.hparams.shift_step, use_cache=False, ) def create_metrics(self): """Creates metrics that use shifters when computing the losses (MSE, MAE, SSIM. TV) for the train, validation and test sets.""" self.metrics_baseline = {} self.metrics_hr = {} # TODO: dict of train/val/test metrics self.train_metrics = self.metrics_factory("train", self.shifter) self.val_metrics = self.metrics_factory("val", self.shifter) self.test_metrics = self.metrics_factory("test", self.shifter) self.baseline_train_metrics = self.metrics_factory( "train", self.baseline_shifter ) self.baseline_val_metrics = self.metrics_factory("val", self.baseline_shifter) self.baseline_test_metrics = self.metrics_factory("test", self.baseline_shifter) # Create metrics for train, val, test def metrics_factory(self, prefix, shifter): """Creates a dictionary of metrics that use the shifter when computing the losses (MSE, MAE, SSIM. TV) for the train, validation and test sets. Parameters ---------- prefix : str Prefix for the metrics (train/val/test). shifter : losses.Shift Shifter to use when computing the losses. Returns ------- dict Dictionary of metrics. """ return { f"{prefix}/MSE": lambda y_hat, y: shifter.registered_loss(mse_loss)( y_hat, y ), f"{prefix}/MAE": lambda y_hat, y: shifter.registered_loss(mae_loss)( y_hat, y ), f"{prefix}/SSIM_loss": lambda y_hat, y,: shifter.registered_loss(ssim_loss)( y_hat, y, window_size=5 ), f"{prefix}/TV": lambda y_hat, y=None: tv_loss(y_hat), } def forward(self, x, y=None): """Forward pass of the model. Parameters ---------- x : Tensor Input tensor (LR) of dimensions (batch_size, channels, height, width). y : Tensor, optional Target tensor (HR) of dimensions (batch_size, channels, height, width). Returns ------- Tensor Output tensor (y_hat) of dimensions (batch_size, channels, height, width). """ return self.backbone(x, y) def validation_log(self, y, m, baseline_m, prefix): log = pd.DataFrame( columns=[ "MSE", "MAE", "SSIM_loss", "TV", "PSNR", "Baseline MSE", "Baseline PSNR", ] ) mse, mae, ssim, tv, psnr = ( m[f"{prefix}/MSE"], m[f"{prefix}/MAE"], m[f"{prefix}/SSIM_loss"], m[f"{prefix}/TV"], m[f"{prefix}/PSNR"], ) baseline_mse, baseline_psnr = ( baseline_m[f"{prefix}/MSE"], baseline_m[f"{prefix}/PSNR_baseline"], ) baseline_m = list(baseline_m.values()) for batch_item in range(y.shape[0]): # batch_size log.loc[self.logging["current_idx"]] = ( mse[batch_item].detach().item(), mae[batch_item].detach().item(), ssim[batch_item].detach().item(), tv[batch_item].detach().item(), psnr[batch_item].detach().item(), baseline_mse[batch_item].detach().item(), baseline_psnr[batch_item].detach().item(), ) self.logging["current_idx"] += 1 self.logging["log"] = pd.concat([self.logging["log"], log]) def loss(self, batch, metrics, baseline_metrics, prefix): """Computes the loss for a batch using the provided metrics. Parameters ---------- batch : dict Dictionary with the batches for LR/HR/etc. metrics : dict of Callable Dictionary with the metrics to use. baseline_metrics : dict of Callable Dictionary with the metrics to use for the baseline. prefix : str Prefix for the metrics (train/val/test). Returns ------- dict Dictionary with the losses. """ batch = self.transform(batch) # GPU/Batched data augmentation x, y = batch["lr"], batch["hr"] # TODO: Run once and store results in a file. x = self.apply_masks_to_x(x, batch) y_hat = self(x, y) y, y_hat = y[:, 0], y_hat[:, 0] y_hat = self.bias_adjust(y, y_hat) self.initialise_shifters() m, mse, mae, ssim, tv = self.compute_metrics(y, y_hat, metrics, prefix) y_hat_base = self.baseline(x)[:, 0] y_hat_base = self.bias_adjust(y, y_hat_base) baseline_m = self.compute_baseline_metrics( y, y_hat_base, m, baseline_metrics, prefix ) loss = ( (self.hparams.w_mse * mse) + (self.hparams.w_mae * mae) + (self.hparams.w_tv * tv) + (self.hparams.w_ssim * ssim) ) if self.hparams.benchmark: self.validation_log(y, detach_dict(m), baseline_m, prefix) return { "y_hat": y_hat, "y_hat_base": y_hat_base, "y": y, "metrics": m, "baseline_metrics": baseline_m, "loss": loss, } def apply_masks_to_x(self, x, batch): """Applies the input masks to the input tensor. Parameters ---------- x : torch.Tensor Input tensor of dimensions (batch_size, channels, height, width). batch : dict Dictionary with the batches for LR/Masks/HR/etc. Returns ------- torch.Tensor Input tensor with the mask channels concatenated. """ if self.hparams.use_masks: masks = batch["lrc"] x = torch.cat([x, masks], dim=-3) # concat channels return x @staticmethod def bias_adjust(y, y_hat): """Applies the brightness correction to the output. Parameters ---------- y : torch.Tensor Target tensor (HR) of dimensions (batch_size, channels, height, width). y_hat : torch.Tensor Output tensor (SR) of dimensions (batch_size, channels, height, width). Returns ------- torch.Tensor Output tensor with the brightness correction applied. """ bias = (y - y_hat).mean(dim=(-1, -2), keepdim=True) # bias / zero-order return y_hat + bias def initialise_shifters(self): """Initialises the shifter for the HR/SR pairs and the HR/LR Baseline pairs.""" self.shifter.y, self.shifter.y_hat = None, None self.baseline_shifter.y, self.baseline_shifter.y_hat = None, None # NOTE: whatever gets cached early on, might be suboptimal. # So it would be good to update the cache periodically. # Clearing the cache every 50 epochs. if self.current_epoch % 50 == 0: self.clear_shifter_cache() def compute_metrics(self, y, y_hat, metrics, prefix): """Computes the metrics for a batch (MSE, MAE, SSIM, TV, PSNR). Parameters ---------- y : torch.Tensor Target tensor (HR) of dimensions (batch_size, channels, height, width). y_hat : torch.Tensor Output tensor (SR) of dimensions (batch_size, channels, height, width). metrics : dict of Callable Dictionary with the metric functions to use. prefix : str Prefix for the metrics (train/val/test). Returns ------- dict, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor Dictionary with the metrics (MSE, MAE, SSIM, TV, PSNR). The metrics themselves, the MSE, MAE, SSIM, TV. """ metrics = { metrics_function: metrics[metrics_function](y_hat, y) for metrics_function in metrics } mse = metrics[f"{prefix}/MSE"].clone() mae = metrics[f"{prefix}/MAE"].clone() ssim = metrics[f"{prefix}/SSIM_loss"].clone() tv = metrics[f"{prefix}/TV"].clone() metrics[f"{prefix}/PSNR"] = -10.0 * torch.log10(mse) return metrics, mse, mae, ssim, tv def compute_baseline_metrics( self, y, y_hat_base, metrics, baseline_metrics, prefix ): """Computes the baseline metrics for a batch (MSE, MAE, SSIM, TV, PSNR_baseline, PSNR_ratio). Parameters ---------- y : torch.Tensor Target tensor (HR) of dimensions (batch_size, channels, height, width). y_hat_base : torch.Tensor Baseline tensor (single upscaled LR) of dimensions (batch_size, channels, height, width). metrics : dict of Callable Dictionary with the metric functions to use. baseline_metrics : dict of Callable Dictionary with the baseline metric functions to use. prefix : str Prefix for the metrics (train/val/test). Returns ------- dict Dictionary with the baseline metrics (MSE, MAE, SSIM, TV, PSNR_baseline, PSNR_ratio). """ baseline_m = {k: baseline_metrics[k](y_hat_base, y) for k in baseline_metrics} baseline_psnr = -10.0 * torch.log10(mse_loss(y_hat_base, y)) baseline_m[f"{prefix}/PSNR_baseline"] = baseline_psnr baseline_m[f"{prefix}/PSNR_ratio"] = ( baseline_m[f"{prefix}/PSNR_baseline"] / metrics[f"{prefix}/PSNR"] ) baseline_m = detach_dict(baseline_m) return baseline_m def clear_shifter_cache(self): """Clears the cache of the output and baseline shifters.""" self.shifter.shift_cache = {} self.baseline_shifter.shift_cache = {} def training_step(self, batch): """Training step. Calls the forward pass, computes the loss and metrics for the training batch. Logs the reduced metrics. Parameters ---------- batch : dict Dictionary with the batches for LR/HR/etc. Returns ------- dict Dictionary with the training loss and metrics. """ loss_output = self.loss( batch, self.train_metrics, self.baseline_train_metrics, prefix="train" ) metrics, metrics_reduced = self.unpack_and_reduce_metrics( loss_output, prefix="train" ) self.log_dict(metrics_reduced, on_step=True, on_epoch=True) return {"loss": metrics_reduced["train/loss"], "metrics": metrics} def unpack_and_reduce_metrics(self, loss_output, prefix): """Unpacks the metrics for a split (prefix) and reduces them (mean). Parameters ---------- loss_output : dict Dictionary with the loss and metrics. prefix : str Prefix for the metrics (train/val/test). Returns ------- dict Dictionary with the reduced (mean) metrics. """ metrics = loss_output["metrics"] metrics_reduced = { metric_name: metric.mean() for metric_name, metric in metrics.items() } metrics_reduced[f"{prefix}/loss"] = loss_output["loss"].mean() return metrics, metrics_reduced def validation_step(self, batch, batch_idx): """Validation step. Calls the forward pass, computes the loss and metrics for the validation batch. Logs the reduced metrics. Parameters ---------- batch : dict Dictionary with the batches for LR/HR/etc. Returns ------- torch.Tensor Validation loss. """ loss_output = self.loss( batch, self.val_metrics, self.baseline_val_metrics, prefix="val" ) _, metrics_reduced = self.unpack_and_reduce_metrics(loss_output, prefix="val") self.log_dict(metrics_reduced) return loss_output["loss"] # Instance-wise loss def validation_epoch_end(self, validation_step_output): """Validation epoch end. Concatenates the validation losses and logs them as a histogram to WandB. Parameters ---------- validation_step_output : dict Dictionary with the validation losses. """ instance_loss = torch.cat(validation_step_output).to("cpu") try: self.logger.experiment.log( { "val/loss_hist": wandb.Histogram(instance_loss), "global_step": self.global_step, } ) except ValueError as e: print(e) def test_step(self, batch, batch_idx): """Test step.+ Calls the forward pass, computes the loss and metrics for the test batch. Logs the reduced metrics. Parameters ---------- batch : dict Dictionary with the batches for LR/HR/etc. Returns ------- torch.Tensor Test loss. """ loss_output = self.loss( batch, self.test_metrics, self.baseline_test_metrics, prefix="test" ) _, metrics_reduced = self.unpack_and_reduce_metrics(loss_output, prefix="test") self.log_dict(metrics_reduced) return loss_output["loss"] # Instance-wise loss def configure_optimizers(self): """Configures the optimizers for the model based on the hyperparameters. Uses the learning rate and weight decay parameters for the Adam optimizer. Uses the learning rate, weight decay and momentum parameters for the SGD optimizer. Uses the cos_anneal_T_0 hyperparameter for the cosine annealing scheduler. See: https://arxiv.org/abs/1608.03983 Returns ------- dict Dictionary containing the optimizer, scheduler and monitor key (val/loss). """ if self.hparams.optimizer == "adam": optimizer = Adam( self.parameters(), lr=self.hparams.learning_rate, weight_decay=self.hparams.weight_decay, ) elif self.hparams.optimizer == "sgd": optimizer = SGD( self.parameters(), lr=self.hparams.learning_rate, weight_decay=self.hparams.weight_decay, momentum=self.hparams.momentum, ) scheduler = CosineAnnealingWarmRestarts( optimizer, T_0=self.hparams.cos_anneal_T_0 ) return { "optimizer": optimizer, "lr_scheduler": scheduler, "monitor": "val/loss", } @staticmethod def add_model_specific_args(parent_parser): """Add model-specific arguments (and their default values) to the parser. Parameters ---------- parent_parser : argparse.ArgumentParser argparse.ArgumentParser to add to. Returns ------- ArgumentParser ArgumentParser with the model-specific arguments added. """ parser = ArgumentParser(parents=[parent_parser], add_help=False) # Model arguments parser.add_argument("--model", default="srcnn", type=str) parser.add_argument("--hidden_channels", default=64, type=int) parser.add_argument("--residual_layers", default=4, type=int) parser.add_argument("--homography_fc_size", default=128, type=int) parser.add_argument("--kernel_size", default=3, type=int) parser.add_argument("--zoom_factor", default=2, type=int) parser.add_argument("--sr_kernel_size", default=1, type=int) # Optimizer arguments parser.add_argument("--optimizer", default="adam", type=str) parser.add_argument("--learning_rate", default=5e-4, type=float) parser.add_argument("--learning_rate_decay", default=0.97, type=float) parser.add_argument("--learning_rate_patience", default=3, type=int) parser.add_argument("--momentum", type=float, default=0.9) parser.add_argument("--weight_decay", type=float, default=1e-4) parser.add_argument("--cos_anneal_T_0", type=int, default=300) # Data arguments parser.add_argument("--input_size", nargs="+", default=(400, 400), type=int) parser.add_argument("--output_size", nargs="+", default=(1000, 1000), type=int) parser.add_argument("--chip_size", nargs="+", type=int, default=(50, 50)) parser.add_argument("--chip_stride", nargs="+", type=int, default=None) parser.add_argument("--revisits", default=8, type=int) parser.add_argument( "--use_masks", action="store_true", help="Include SCL cloud masks as extra input bands.", ) # Shifter arguments parser.add_argument("--registration_kind", default=None, type=str) parser.add_argument("--shift_px", default=2, type=int) parser.add_argument("--shift_mode", default="lanczos", type=str) parser.add_argument("--shift_step", default=0.5, type=float) parser.add_argument( "--use_reference_frame", action="store_true", help="Pair each revisit with a reference (median) frame.", ) # Loss weights: MSE, MAE, SSIM, TV parser.add_argument("--w_mse", default=1.0, type=float) parser.add_argument("--w_mae", default=0.0, type=float) parser.add_argument("--w_ssim", default=0.0, type=float) parser.add_argument("--w_tv", default=0.0, type=float) return parser class ImagePredictionLogger(pl.Callback): """Logs the model losses, inputs and outputs to WandB. Parameters ---------- pl : pl.LightningModule Lightning module containing the model. """ def __init__( self, train_dataloader, val_dataloader, test_dataloader, log_every_n_epochs=1, window_size=None, ): """Initialize the logger. Parameters ---------- train_dataloader : torch.utils.data.DataLoader Training dataloader. val_dataloader : torch.utils.data.DataLoader Validation dataloader. test_dataloader : torch.utils.data.DataLoader Test dataloader. log_every_n_epochs : int, optional Log every n epochs. window_size : tuple of int, optional Size of the window to use for the patches. """ super().__init__() self.train_dataloader = train_dataloader self.val_dataloader = val_dataloader self.test_dataloader = test_dataloader self.log_every_n_epochs = log_every_n_epochs self.window_size = window_size self.baseline_psnr = None def _on_epoch_end(self, prefix, batch, trainer, pl_module): """Log the current epoch's metrics and images. Parameters ---------- prefix : str The split prefix (train, val, test). batch : dict A batch of data. trainer : pytorch_lightning.trainer.Trainer The trainer. pl_module : pytorch_lightning.LightningModule The model. """ lr, hr = batch["lr"], batch["hr"] batch = self.move_batch_to_device(batch, pl_module.device) metrics, baseline_metrics = self.fetch_metrics_for_split(prefix, pl_module) loss_output, instance_psnr, baseline_psnr = self.get_losses_from_metrics( batch, metrics, baseline_metrics, pl_module, prefix ) sr, sr_base = self.get_output_and_baseline(loss_output) window_size, window_size_hr = self.calculate_window_size(lr, hr, pl_module) lr, hr, sr, sr_base = self.extract_patches_from_images( sr, sr_base, lr, hr, window_size, window_size_hr ) instance_psnr, baseline_psnr = self.extract_psnr_values_for_patches( sr, sr_base, lr, instance_psnr, baseline_psnr ) lr = self.convert_lr_to_RGB_only_for_logging(lr) self.log_image_previews_to_wandb( sr, sr_base, lr, hr, instance_psnr, baseline_psnr, pl_module, prefix, trainer, ) def move_batch_to_device(self, batch, device): """Move all tensors from batch to the specified device. Parameters ---------- batch : dict A batch of data. device : torch.device The device to move the tensors to. Returns ------- dict The batch with all tensors moved to the specified device. """ return {key: value.to(device) for key, value in batch.items()} def fetch_metrics_for_split(self, prefix, pl_module): """Fetch the metrics for the split. Parameters ---------- prefix : str The prefix for the metrics. pl_module : pytorch_lightning.LightningModule The model. Returns ------- dict, dict The metrics and baseline metrics for the current batch. """ if prefix == "val": metrics = pl_module.val_metrics baseline_metrics = pl_module.baseline_val_metrics elif prefix == "train": metrics = pl_module.train_metrics baseline_metrics = pl_module.baseline_train_metrics elif prefix == "test": metrics = pl_module.test_metrics baseline_metrics = pl_module.baseline_test_metrics return metrics, baseline_metrics def get_losses_from_metrics( self, batch, metrics, baseline_metrics, pl_module, prefix ): """Get the losses that will be logged from the metrics. Parameters ---------- batch : dict The batch. metrics : dict The metrics for the current batch. baseline_metrics : dict The baseline metrics for the current batch. pl_module : pytorch_lightning.LightningModule The model. prefix : str The prefix for the metrics. Returns ------- torch.Tensor, torch.Tensor, torch.Tensor The loss output, the PSNR values for the SR and Baseline images in the batch. """ loss_output = pl_module.loss(batch, metrics, baseline_metrics, prefix=prefix) instance_psnr = loss_output["metrics"][f"{prefix}/PSNR"].to("cpu") baseline_psnr = loss_output["baseline_metrics"][f"{prefix}/PSNR_baseline"].to( "cpu" ) return loss_output, instance_psnr, baseline_psnr def get_output_and_baseline(self, loss_output): """Get the output and baseline images from the loss output. Parameters ---------- loss_output : dict The loss output. Returns ------- torch.Tensor, torch.Tensor The output and baseline images. """ sr = loss_output["y_hat"].to("cpu")[:, None] sr_base = loss_output["y_hat_base"].to("cpu")[:, None] return sr, sr_base def calculate_window_size(self, lr, hr, pl_module): """Calculate the window size for the patches. Parameters ---------- lr : torch.Tensor The LR image. hr : torch.Tensor The HR image. pl_module : pytorch_lightning.LightningModule The model. Returns ------- tuple of int, tuple of int The window size for lr and hr. """ _, _, _, lr_height, lr_width = lr.shape hr_height, hr_width = hr.shape[-2:] ratio_hr_lr_height, ratio_hr_lr_width = ( hr_height / lr_height, hr_width / lr_width, ) window_size = self.window_size or pl_module.hparams.input_size assert window_size <= tuple(pl_module.hparams.chip_size), "window_size too big" window_size_hr = ( round(ratio_hr_lr_height * window_size[0]), round(ratio_hr_lr_width * window_size[0]), ) return window_size, window_size_hr def extract_patches_from_images( self, sr, sr_base, lr, hr, window_size, window_size_hr ): """Extract the patches from the images. Parameters ---------- sr : torch.Tensor The SR image. sr_base : torch.Tensor The baseline SR image. hr : torch.Tensor The HR image. window_size : tuple of int The window size. window_size_hr : tuple of int The window size for the HR image. Returns ------- torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor The extracted SR, SR_base, LR, HR images. """ lr = self.batch_patch_extractor( lr, window_size ) # (batch*patches, tensors/revisits, channels, patch_height, patch_width) hr = self.batch_patch_extractor(hr, window_size_hr) sr = self.batch_patch_extractor(sr, window_size_hr) sr_base = self.batch_patch_extractor(sr_base, window_size_hr) return lr, hr, sr, sr_base @staticmethod def batch_patch_extractor(x, window_size): """Extract patches from a batch of images. Parameters ---------- x : Tensor A batch of images (batch, tensors/revisits, channels, height, width). window_size : tuple of int The size of the patch to extract. Returns ------- Tensor A batch of patches extracted from the images (batch*patches, tensors, channels, patch_height, patch_width), with patches extracted patches. """ batch, tensors, channels, height, width = x.shape x = extract_tensor_patches( x.view(batch, tensors * channels, height, width), window_size, stride=window_size, ) _, patches, _, patch_height, patch_width = x.shape return x.view( batch, patches, tensors, channels, patch_height, patch_width ).view(batch * patches, tensors, channels, patch_height, patch_width) def extract_psnr_values_for_patches( self, sr, sr_base, lr, instance_psnr, baseline_psnr ): """Extract the PSNR values for the patches. Parameters ---------- sr : torch.Tensor The SR image. sr_base : torch.Tensor The baseline SR image. lr : torch.Tensor The LR image. instance_psnr : torch.Tensor The instance PSNR values. baseline_psnr : torch.Tensor The baseline PSNR values. Returns ------- torch.Tensor, torch.Tensor The extracted instance and baseline PSNR values. """ instance_psnr = ( instance_psnr[None].expand(sr.shape[0] // lr.shape[0], -1).flatten() ) baseline_psnr = ( baseline_psnr[None].expand(sr_base.shape[0] // lr.shape[0], -1).flatten() ) return instance_psnr, baseline_psnr def convert_lr_to_RGB_only_for_logging(self, lr): """Convert the LR images to RGB only for logging. Parameters ---------- lr : torch.Tensor The LR image. Returns ------- torch.Tensor The LR image in RGB. """ if ( lr.shape[2] > 3 ): # If all bands are used, get only the RGB bands for WandB image logging lr = lr[:, :, S2_ALL_12BANDS["true_color_zero_index"]] return lr def log_image_previews_to_wandb( self, sr, sr_base, lr, hr, instance_psnr, baseline_psnr, pl_module, prefix, trainer, ): """Log the current epoch's images to WandB. Parameters ---------- sr : torch.Tensor The SR image. sr_base : torch.Tensor The baseline SR image. lr : torch.Tensor The LR image. hr : torch.Tensor The HR image. instance_psnr : torch.Tensor The instance PSNR values. baseline_psnr : torch.Tensor The baseline PSNR values. pl_module : pl.LightningModule The model. prefix : str The split prefix (train, val, test). trainer : pytorch_lightning.trainer.Trainer The trainer. """ number_of_rows_for_lr_revisits = int(sqrt(pl_module.hparams.revisits).round()) trainer.logger.experiment.log( { f"{prefix}-lr": [ wandb.Image( make_grid( lr_image, nrow=number_of_rows_for_lr_revisits, normalize=True, scale_each=True, ), caption=f"lr", ) for lr_image in lr ], f"{prefix}-hr": [ wandb.Image( make_grid(hr_image, nrow=1, normalize=True, scale_each=True), caption=f"hr", ) for hr_image in hr ], f"{prefix}-baseline_upscaled_bicubic": [ wandb.Image( make_grid( baseline_image, nrow=1, normalize=True, scale_each=True ), caption=f"baseline:{str(baseline_psnr)}", ) for baseline_image, baseline_psnr in zip(sr_base, baseline_psnr) ], f"{prefix}-sr": [ wandb.Image( make_grid(image, nrow=1, normalize=True, scale_each=True), caption=f"sr:{str(sr_psnr)}", ) for image, sr_psnr in zip(sr, instance_psnr) ], } ) def on_validation_epoch_end(self, trainer, pl_module): """Called at the end of the validation epoch. If the current epoch should be logged (i.e. the current epoch is a multiple of the log interval), then _on_epoch_end is called for the next batch of the validation set. Parameters ---------- trainer : pytorch_lightning.trainer.Trainer The trainer. pl_module : pytorch_lightning.LightningModule The model. """ if trainer.current_epoch % self.log_every_n_epochs == 0: batch = next(iter(self.val_dataloader)) self._on_epoch_end("val", batch, trainer, pl_module) def on_test_epoch_end(self, trainer, pl_module): """Called at the end of the test epoch. Parameters ---------- trainer : pytorch_lightning.trainer.Trainer The trainer. pl_module : pytorch_lightning.LightningModule The model. """ for batch in iter(self.test_dataloader): self._on_epoch_end("test", batch, trainer, pl_module)