from dataclasses import dataclass from pathlib import Path from typing import Literal, Optional import torch from einops import rearrange from lightning.pytorch import LightningModule from lightning.pytorch.utilities import rank_zero_only from tabulate import tabulate from torch import Tensor, nn import torch.nn.functional as F from ..dataset.data_module import get_data_shim from ..dataset.types import BatchedExample from ..evaluation.metrics import ( compute_lpips, compute_psnr, compute_ssim, ) from ..global_cfg import get_cfg from ..loss import Loss from ..misc.benchmarker import Benchmarker from ..misc.image_io import prep_image, save_image from ..misc.step_tracker import StepTracker from ..misc.utils import ( get_rank, inverse_normalize, vis_depth_map, ) from ..visualization.annotation import add_label from ..visualization.layout import add_border, hcat, vcat from .decoder.decoder import DepthRenderingMode from lightning.pytorch.loggers.wandb import WandbLogger @dataclass class OptimizerCfg: lr: float warm_up_steps: int backbone_lr_multiplier: float @dataclass class TestCfg: output_path: Path align_pose: bool pose_align_steps: int rot_opt_lr: float trans_opt_lr: float compute_scores: bool save_image: bool save_video: bool save_compare: bool generate_video: bool mode: Literal["inference", "evaluation"] image_folder: str @dataclass class TrainCfg: output_path: Path depth_mode: DepthRenderingMode | None extended_visualization: bool print_log_every_n_steps: int distiller: str distill_max_steps: int pose_loss_alpha: float = 1.0 pose_loss_delta: float = 1.0 cxt_depth_weight: float = 0.01 weight_pose: float = 1.0 weight_depth: float = 1.0 weight_normal: float = 1.0 render_ba: bool = False render_ba_after_step: int = 0 class ModelWrapper(LightningModule): logger: Optional[WandbLogger] model: nn.Module losses: nn.ModuleList optimizer_cfg: OptimizerCfg test_cfg: TestCfg train_cfg: TrainCfg step_tracker: StepTracker | None def __init__( self, optimizer_cfg: OptimizerCfg, test_cfg: TestCfg, train_cfg: TrainCfg, model: nn.Module, losses: list[Loss], step_tracker: StepTracker | None, ) -> None: super().__init__() self.optimizer_cfg = optimizer_cfg self.test_cfg = test_cfg self.train_cfg = train_cfg self.step_tracker = step_tracker # Set up the model. self.encoder_visualizer = None self.model = model self.data_shim = get_data_shim(self.model.encoder) self.losses = nn.ModuleList(losses) # This is used for testing. self.benchmarker = Benchmarker() @staticmethod def flatten_indices(indices) -> list[int]: if isinstance(indices, Tensor): return [int(idx) for idx in indices.detach().cpu().reshape(-1).tolist()] if isinstance(indices, (list, tuple)): flattened = [] for item in indices: flattened.extend(ModelWrapper.flatten_indices(item)) return flattened return [int(indices)] def on_train_epoch_start(self) -> None: # our custom dataset and sampler has to have epoch set by calling set_epoch print(f"Train epoch start on rank {self.trainer.global_rank}") if hasattr(self.trainer.datamodule.train_loader.dataset, "set_epoch"): self.trainer.datamodule.train_loader.dataset.set_epoch(self.current_epoch) if hasattr(self.trainer.datamodule.train_loader.sampler, "set_epoch"): self.trainer.datamodule.train_loader.sampler.set_epoch(self.current_epoch) def on_validation_epoch_start(self) -> None: print(f"Validation epoch start on rank {self.trainer.global_rank}") # our custom dataset and sampler has to have epoch set by calling set_epoch if hasattr(self.trainer.datamodule.val_loader.dataset, "set_epoch"): self.trainer.datamodule.val_loader.dataset.set_epoch(self.current_epoch) if hasattr(self.trainer.datamodule.val_loader.sampler, "set_epoch"): self.trainer.datamodule.val_loader.sampler.set_epoch(self.current_epoch) def training_step(self, batch, batch_idx): if isinstance(batch, list): batch_combined = None for batch_per_dl in batch: if batch_combined is None: batch_combined = batch_per_dl else: for k in batch_combined.keys(): if isinstance(batch_combined[k], list): batch_combined[k] += batch_per_dl[k] elif isinstance(batch_combined[k], dict): for kk in batch_combined[k].keys(): batch_combined[k][kk] = torch.cat( [batch_combined[k][kk], batch_per_dl[k][kk]], dim=0 ) else: raise NotImplementedError batch = batch_combined batch: BatchedExample = self.data_shim(batch) context_image = (batch["context"]["image"] + 1) / 2 # Run the model. encoder_output, output = self.model(context_image, self.global_step) gaussians, pred_pose_enc_list, depth_dict = ( encoder_output.gaussians, encoder_output.pred_pose_enc_list, encoder_output.depth_dict, ) distill_infos = encoder_output.distill_infos target_gt = (batch["context"]["image"] + 1) / 2 num_context_views = target_gt.shape[1] using_index = torch.arange(num_context_views, device=gaussians.means.device) batch["using_index"] = using_index psnr_probabilistic = compute_psnr( rearrange(target_gt, "b v c h w -> (b v) c h w"), rearrange(output.color, "b v c h w -> (b v) c h w"), ) self.log("train/psnr_probabilistic", psnr_probabilistic.mean().item()) total_loss = 0 with torch.amp.autocast("cuda", enabled=False): depth_loss_idx = list(get_cfg()["loss"].keys()).index("depth") depth_loss_fn = self.losses[depth_loss_idx].ctx_depth_loss loss_depth_ctx = depth_loss_fn( depth_dict["depth"], batch, cxt_depth_weight=self.train_cfg.cxt_depth_weight, ) self.log("loss/loss_depth_ctx", loss_depth_ctx.item()) total_loss = total_loss + loss_depth_ctx for loss_fn in self.losses: if loss_fn.name == "depth": break loss = loss_fn.forward( output, batch, gaussians, depth_dict, self.global_step ) self.log(f"loss/{loss_fn.name}", loss.item()) total_loss = total_loss + loss loss_ca1 = F.mse_loss( pred_pose_enc_list, distill_infos["pred_pose_enc_list"][:, 0:1, -1] ) loss_ca = 10 * loss_ca1 self.log("loss/loss_ca", loss_ca.item()) total_loss = total_loss + loss_ca self.log("loss/total", total_loss.item()) self.log("info/global_step", self.global_step) if self.step_tracker is not None: self.step_tracker.set_step(self.global_step) del batch return total_loss def on_after_backward(self): for name, p in self.named_parameters(): if p.grad is None: continue grad = p.grad.detach() if torch.isnan(grad).any() or torch.isinf(grad).any(): print(f"[NaN-Guard]") p.grad = torch.zeros_like(p.grad) continue @rank_zero_only def validation_step(self, batch, batch_idx, dataloader_idx=0): batch: BatchedExample = self.data_shim(batch) total_batches = len(self.trainer.datamodule.val_loader) print(f"Rank {self.global_rank}, batch {batch_idx+1}/{total_batches}") print( f"validation step {self.global_step}; " f"scene = {batch['scene']}; " f"context = {batch['context']['index'].tolist()}" ) # Render Gaussians. b, v, _, h, w = batch["context"]["image"].shape assert b == 1 encoder_output, output = self.model( (batch["context"]["image"] + 1) / 2, self.global_step, ) # Compute validation metrics. rgb_pred = output.color[0].float() rgb_gt = (batch["context"]["image"][0].float() + 1) / 2 psnr = compute_psnr(rgb_gt, rgb_pred).mean() self.log(f"val/psnr", psnr) lpips = compute_lpips(rgb_gt, rgb_pred).mean() self.log(f"val/lpips", lpips) ssim = compute_ssim(rgb_gt, rgb_pred).mean() self.log(f"val/ssim", ssim) # Construct comparison image. context_img = inverse_normalize(batch["context"]["image"][0]) context = [] for i in range(context_img.shape[0]): context.append(context_img[i]) depth_dict = encoder_output.depth_dict model_depth_pred = depth_dict["depth"].squeeze(-1)[0] model_depth_pred = vis_depth_map(model_depth_pred) depth_pred = vis_depth_map(output.depth[0]) comparison = hcat( add_label(vcat(*context), "Context"), add_label(vcat(*rgb_gt), "Target (Ground Truth)"), add_label(vcat(*rgb_pred), "Rendered Target"), add_label(vcat(*depth_pred), "Rendered Depth"), add_label(vcat(*model_depth_pred), "GS Depth"), ) comparison = torch.nn.functional.interpolate( comparison.unsqueeze(0), scale_factor=0.5, mode="bicubic", align_corners=False, ).squeeze(0) self.logger.log_image( "comparison", [prep_image(add_border(comparison))], step=self.global_step, caption=batch["scene"], ) if self.encoder_visualizer is not None: for k, image in self.encoder_visualizer.visualize( batch["context"], self.global_step ).items(): self.logger.log_image(k, [prep_image(image)], step=self.global_step) def test_step(self, batch, batch_idx): batch: BatchedExample = self.data_shim(batch) b, v, _, h, w = batch["target"]["image"].shape assert b == 1 if batch_idx % 100 == 0: print( f"Rank {get_rank()} test step {batch_idx:0>6}; " f"scene = {batch['scene']}; " f"context = {self.flatten_indices(batch['context']['index'])}; " f"target = {self.flatten_indices(batch['target']['index'])}" ) # Render Gaussians. with torch.no_grad(): with self.benchmarker.time("encoder"): ( gaussians, pred_all_extrinsic, pred_context_pose, ) = self.model.encoder.inference( (batch["context"]["image"] + 1) / 2, (batch["target"]["image"] + 1) / 2, global_step=self.global_step, ) num_context_view = batch["context"]["image"].shape[1] pred_all_context_extrinsic, pred_all_target_extrinsic = ( pred_all_extrinsic[:, :num_context_view], pred_all_extrinsic[:, num_context_view:], ) scale_factor = ( pred_context_pose["extrinsic"][:, :, :3, 3].mean() / pred_all_context_extrinsic[:, :, :3, 3].mean() ) pred_all_target_extrinsic[..., :3, 3] = ( pred_all_target_extrinsic[..., :3, 3] * scale_factor ) pred_all_context_extrinsic[..., :3, 3] = ( pred_all_context_extrinsic[..., :3, 3] * scale_factor ) with self.benchmarker.time("decoder", num_calls=v): output = self.model.decoder.forward( gaussians, pred_all_target_extrinsic, pred_context_pose["intrinsic"][:, 0:1, :, :] .repeat(1, pred_all_target_extrinsic.shape[1], 1, 1) .float(), torch.ones(1, v, device="cuda") * 0.01, torch.ones(1, v, device="cuda") * 100, (h, w), ) psnr = None with torch.no_grad(): if self.test_cfg.compute_scores: rgb_pred = output.color[0] rgb_gt = batch["target"]["image"][0] psnr = compute_psnr(rgb_gt, rgb_pred).mean().item() all_metrics = { f"lpips_ours": compute_lpips(rgb_gt, rgb_pred).mean().item(), f"ssim_ours": compute_ssim(rgb_gt, rgb_pred).mean().item(), f"psnr_ours": psnr, } methods = ["ours"] self.log_dict(all_metrics, prog_bar=True, sync_dist=True, on_epoch=True) self.print_preview_metrics(all_metrics, methods) # Save images. (scene,) = batch["scene"] name = get_cfg()["wandb"]["name"] path = self.test_cfg.output_path / name scene_dir = f"{psnr:.4f}_{scene}" if psnr is not None else scene target_indices = self.flatten_indices(batch["target"]["index"]) for i, idx in enumerate(target_indices): single_color = output.color[0][i] res_path = path / scene_dir / "color" / f"{int(idx):0>6}.png" save_image(single_color, res_path) def on_test_end(self) -> None: self.benchmarker.summarize() def print_preview_metrics( self, metrics: dict[str, float | Tensor], methods: list[str] | None = None, overlap_tag: str | None = None, ) -> None: if getattr(self, "running_metrics", None) is None: self.running_metrics = metrics self.running_metric_steps = 1 else: s = self.running_metric_steps self.running_metrics = { k: ((s * v) + metrics[k]) / (s + 1) for k, v in self.running_metrics.items() } self.running_metric_steps += 1 if overlap_tag is not None: if getattr(self, "running_metrics_sub", None) is None: self.running_metrics_sub = {overlap_tag: metrics} self.running_metric_steps_sub = {overlap_tag: 1} elif overlap_tag not in self.running_metrics_sub: self.running_metrics_sub[overlap_tag] = metrics self.running_metric_steps_sub[overlap_tag] = 1 else: s = self.running_metric_steps_sub[overlap_tag] self.running_metrics_sub[overlap_tag] = { k: ((s * v) + metrics[k]) / (s + 1) for k, v in self.running_metrics_sub[overlap_tag].items() } self.running_metric_steps_sub[overlap_tag] += 1 metric_list = ["psnr", "lpips", "ssim"] def print_metrics(runing_metric, methods=None): table = [] if methods is None: methods = ["ours"] for method in methods: row = [ f"{runing_metric[f'{metric}_{method}']:.3f}" for metric in metric_list ] table.append((method, *row)) headers = ["Method"] + metric_list table = tabulate(table, headers) print(table) print("All Pairs:") print_metrics(self.running_metrics, methods) def configure_optimizers(self): new_params, new_param_names = [], [] for name, param in self.named_parameters(): if not param.requires_grad: continue new_params.append(param) new_param_names.append(name) param_dicts = [ { "params": new_params, "lr": self.optimizer_cfg.lr, } ] optimizer = torch.optim.AdamW( param_dicts, lr=self.optimizer_cfg.lr, weight_decay=0.1, betas=(0.9, 0.95) ) max_steps = get_cfg()["trainer"]["max_steps"] warm_up_steps = self.optimizer_cfg.warm_up_steps if warm_up_steps > 0: warm_up = torch.optim.lr_scheduler.LinearLR( optimizer, start_factor=1.0 / warm_up_steps, end_factor=1.0, total_iters=warm_up_steps, ) lr_scheduler_cosine = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=max_steps - warm_up_steps, eta_min=self.optimizer_cfg.lr * 0.1, ) lr_scheduler = torch.optim.lr_scheduler.SequentialLR( optimizer, schedulers=[warm_up, lr_scheduler_cosine], milestones=[warm_up_steps], ) else: lr_scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=max_steps, eta_min=self.optimizer_cfg.lr * 0.1 ) return { "optimizer": optimizer, "lr_scheduler": { "scheduler": lr_scheduler, "interval": "step", "frequency": 1, }, }