diff --git a/nerfactory/models/base.py b/nerfactory/models/base.py index 5067c840c..d13ce3a3a 100644 --- a/nerfactory/models/base.py +++ b/nerfactory/models/base.py @@ -20,7 +20,7 @@ from __future__ import annotations from abc import abstractmethod from collections import defaultdict -from typing import Any, Dict, List, Optional, Tuple, Union, overload +from typing import Any, Dict, List, Optional, Tuple import torch from torch import nn @@ -115,21 +115,9 @@ class Model(nn.Module): Outputs of model. (ie. rendered colors) """ - @overload - def forward(self, ray_bundle: RayBundle, batch: None = None) -> Dict[str, torch.Tensor]: - ... - - @overload - def forward( - self, ray_bundle: RayBundle, batch: Dict[str, torch.Tensor] - ) -> Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]]: - ... - def forward( self, ray_bundle: RayBundle, batch: Optional[Dict[str, torch.Tensor]] = None - ) -> Union[ - Dict[str, torch.Tensor], Tuple[Dict[str, torch.Tensor], Dict[str, torch.Tensor], Dict[str, torch.Tensor]] - ]: + ) -> Dict[str, torch.Tensor]: """Run forward starting with a ray bundle. This outputs different things depending on the configuration of the model and whether or not the batch is provided (whether or not we are training basically) @@ -157,11 +145,7 @@ class Model(nn.Module): batch = get_masked_dict(batch, valid_mask) # NOTE(ethan): this is really slow if on CPU! outputs = self.get_outputs(intersected_ray_bundle) - metrics_dict = self.get_metrics_dict(outputs=outputs, batch=batch) - loss_dict = self.get_loss_dict( - outputs=outputs, batch=batch, metrics_dict=metrics_dict, loss_coefficients=self.config.loss_coefficients - ) - return outputs, loss_dict, metrics_dict + return outputs def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.Tensor]: """Compute and returns metrics. @@ -175,14 +159,13 @@ class Model(nn.Module): return {} @abstractmethod - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.Tensor]: + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: """Computes and returns the losses dict. Args: outputs: the output to compute loss dict to batch: ground truth batch corresponding to outputs - metrics_dict: collection of metrics to compute - loss_coefficients: list of loss coefficients/weightings to apply + metrics_dict: dictionary of metrics, some of which we can use for loss """ @torch.no_grad() diff --git a/nerfactory/models/compound.py b/nerfactory/models/compound.py index 25f632cd3..8110ad02c 100644 --- a/nerfactory/models/compound.py +++ b/nerfactory/models/compound.py @@ -190,7 +190,7 @@ class CompoundModel(Model): } return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): + def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(self.device) mask = outputs["alive_ray_mask"] rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) diff --git a/nerfactory/models/instant_ngp.py b/nerfactory/models/instant_ngp.py index d85cadf37..4e36be1e0 100644 --- a/nerfactory/models/instant_ngp.py +++ b/nerfactory/models/instant_ngp.py @@ -177,7 +177,7 @@ class NGPModel(Model): } return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): + def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(self.device) mask = outputs["alive_ray_mask"] rgb_loss = self.rgb_loss(image[mask], outputs["rgb"][mask]) diff --git a/nerfactory/models/mipnerf.py b/nerfactory/models/mipnerf.py index 51dc6c8f2..876145564 100644 --- a/nerfactory/models/mipnerf.py +++ b/nerfactory/models/mipnerf.py @@ -137,12 +137,12 @@ class MipNerfModel(Model): } return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): + def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"] rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} - loss_dict = misc.scale_dict(loss_dict, loss_coefficients) + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict def get_image_metrics_and_images( diff --git a/nerfactory/models/mipnerf_360.py b/nerfactory/models/mipnerf_360.py index 89812167d..e49b548db 100644 --- a/nerfactory/models/mipnerf_360.py +++ b/nerfactory/models/mipnerf_360.py @@ -147,7 +147,7 @@ class MipNerf360Model(Model): } return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.Tensor]: + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: image = batch["image"] rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) @@ -157,7 +157,7 @@ class MipNerf360Model(Model): "ray_loss_coarse": torch.mean(outputs["ray_loss_coarse"]), "ray_loss_fine": torch.mean(outputs["ray_loss_fine"]), } - loss_dict = misc.scale_dict(loss_dict, loss_coefficients) + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict def get_image_metrics_and_images( diff --git a/nerfactory/models/nerfw.py b/nerfactory/models/nerfw.py index b6d49e19b..bbb97775b 100644 --- a/nerfactory/models/nerfw.py +++ b/nerfactory/models/nerfw.py @@ -195,7 +195,7 @@ class NerfWModel(Model): } return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): + def get_loss_dict(self, outputs, batch, metrics_dict=None): device = outputs["rgb_coarse"].device image = batch["image"].to(device) rgb_coarse = outputs["rgb_coarse"] @@ -213,7 +213,7 @@ class NerfWModel(Model): "uncertainty_loss": uncertainty_loss, "density_loss": density_loss, } - loss_dict = misc.scale_dict(loss_dict, loss_coefficients) + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict def get_image_metrics_and_images( diff --git a/nerfactory/models/semantic_nerf.py b/nerfactory/models/semantic_nerf.py index b90a12cdd..786d750a0 100644 --- a/nerfactory/models/semantic_nerf.py +++ b/nerfactory/models/semantic_nerf.py @@ -115,7 +115,7 @@ class SemanticNerfModel(NeRFModel): } return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): + def get_loss_dict(self, outputs, batch, metrics_dict=None): image = batch["image"].to(outputs["rgb_coarse"].device) rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) @@ -127,7 +127,7 @@ class SemanticNerfModel(NeRFModel): "rgb_loss_fine": rgb_loss_fine, "semantic_loss_fine": semantic_loss_fine, } - loss_dict = misc.scale_dict(loss_dict, loss_coefficients) + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict def get_image_metrics_and_images( diff --git a/nerfactory/models/tensorf.py b/nerfactory/models/tensorf.py index c739e2d59..e4025fd5e 100644 --- a/nerfactory/models/tensorf.py +++ b/nerfactory/models/tensorf.py @@ -178,7 +178,7 @@ class TensoRFModel(Model): outputs = {"rgb": rgb, "accumulation": accumulation, "depth": depth} return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.Tensor]: + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: # Scaling metrics by coefficients to create the losses. device = outputs["rgb"].device image = batch["image"].to(device) @@ -192,7 +192,7 @@ class TensoRFModel(Model): line_feature_loss = self.feature_loss(line_coef, torch.zeros_like(line_coef)) loss_dict = {"rgb_loss": rgb_loss, "feature_loss": plane_feature_loss + line_feature_loss} - loss_dict = misc.scale_dict(loss_dict, loss_coefficients) + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict def get_image_metrics_and_images( diff --git a/nerfactory/models/vanilla_nerf.py b/nerfactory/models/vanilla_nerf.py index 4693c797e..c5c0cfb72 100644 --- a/nerfactory/models/vanilla_nerf.py +++ b/nerfactory/models/vanilla_nerf.py @@ -166,7 +166,7 @@ class NeRFModel(Model): } return outputs - def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.Tensor]: + def get_loss_dict(self, outputs, batch, metrics_dict=None) -> Dict[str, torch.Tensor]: # Scaling metrics by coefficients to create the losses. device = outputs["rgb_coarse"].device image = batch["image"].to(device) @@ -175,7 +175,7 @@ class NeRFModel(Model): rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) loss_dict = {"rgb_loss_coarse": rgb_loss_coarse, "rgb_loss_fine": rgb_loss_fine} - loss_dict = misc.scale_dict(loss_dict, loss_coefficients) + loss_dict = misc.scale_dict(loss_dict, self.config.loss_coefficients) return loss_dict def get_image_metrics_and_images( diff --git a/nerfactory/pipelines/base.py b/nerfactory/pipelines/base.py index 10183f82f..29d4f209f 100644 --- a/nerfactory/pipelines/base.py +++ b/nerfactory/pipelines/base.py @@ -129,7 +129,10 @@ class Pipeline(nn.Module): assert self.datamanager.train_sampler is not None self.datamanager.train_sampler.set_epoch(step) ray_bundle, batch = self.datamanager.next_train(step) - model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch) + model_outputs = self.model(ray_bundle, batch) + metrics_dict = self.model.get_metrics_dict(model_outputs, batch) + loss_dict = self.model.get_loss_dict(model_outputs, batch, metrics_dict) + return model_outputs, loss_dict, metrics_dict @abstractmethod @@ -146,7 +149,9 @@ class Pipeline(nn.Module): assert self.datamanager.eval_sampler is not None self.datamanager.eval_sampler.set_epoch(step) ray_bundle, batch = self.datamanager.next_eval(step) - model_outputs, loss_dict, metrics_dict = self.model(ray_bundle, batch) + model_outputs = self.model(ray_bundle, batch) + loss_dict = self.model.get_loss_dict(model_outputs, batch) + metrics_dict = self.model.get_metrics_dict(model_outputs, batch) self.train() return model_outputs, loss_dict, metrics_dict