diff --git a/nerfactory/graphs/base.py b/nerfactory/graphs/base.py index a163af624..130f65966 100644 --- a/nerfactory/graphs/base.py +++ b/nerfactory/graphs/base.py @@ -186,12 +186,9 @@ class Graph(AbstractGraph): masked_batch = get_masked_dict(batch, valid_mask) # NOTE(ethan): this is really slow if on CPU! outputs = self.get_outputs(masked_intersected_ray_bundle) metrics_dict = self.get_metrics_dict(outputs=outputs, batch=masked_batch) - loss_dict = self.get_loss_dict(outputs=outputs, batch=masked_batch) - - # scaling losses by coefficients. - for loss_name in loss_dict.keys(): - if loss_name in self.loss_coefficients: - loss_dict[loss_name] *= self.loss_coefficients[loss_name] + loss_dict = self.get_loss_dict( + outputs=outputs, batch=masked_batch, metrics_dict=metrics_dict, loss_coefficients=self.loss_coefficients + ) return outputs, loss_dict, metrics_dict def forward(self, ray_indices: TensorType["num_rays", 3], batch: Union[str, Dict[str, torch.tensor]] = None): @@ -199,16 +196,16 @@ class Graph(AbstractGraph): ray_bundle = self.ray_generator.forward(ray_indices) return self.forward_after_ray_generator(ray_bundle, batch=batch) - @abstractmethod - def get_loss_dict(self, outputs, batch) -> Dict[str, torch.tensor]: - """Computes and returns the losses dict.""" - def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.tensor]: - """Compute and obtain metrics and coefficients.""" + """Compute and returns metrics.""" # pylint: disable=unused-argument # pylint: disable=no-self-use return {} + @abstractmethod + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.tensor]: + """Computes and returns the losses dict.""" + @torch.no_grad() def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle): """Takes in camera parameters and computes the output of the graph.""" diff --git a/nerfactory/graphs/instant_ngp.py b/nerfactory/graphs/instant_ngp.py index a010847fa..6057fdd16 100644 --- a/nerfactory/graphs/instant_ngp.py +++ b/nerfactory/graphs/instant_ngp.py @@ -32,7 +32,7 @@ from nerfactory.fields.modules.field_heads import FieldHeadNames from nerfactory.graphs.base import Graph from nerfactory.graphs.modules.ray_sampler import NGPSpacedSampler from nerfactory.optimizers.loss import MSELoss -from nerfactory.utils import colors, visualization, writer +from nerfactory.utils import colors, misc, visualization, writer from nerfactory.utils.callbacks import Callback @@ -106,7 +106,7 @@ class NGPGraph(Graph): } return outputs - def get_loss_dict(self, outputs, batch): + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): device = self.get_device() image = batch["image"].to(device) if "mask" in outputs: @@ -115,6 +115,7 @@ class NGPGraph(Graph): else: rgb_loss = self.rgb_loss(image, outputs["rgb"]) loss_dict = {"rgb_loss": rgb_loss} + loss_dict = misc.scale_dict(loss_dict, loss_coefficients) return loss_dict def process_outputs_as_images(self, outputs): # pylint:disable=no-self-use diff --git a/nerfactory/graphs/mipnerf.py b/nerfactory/graphs/mipnerf.py index e3535be5a..68fb07683 100644 --- a/nerfactory/graphs/mipnerf.py +++ b/nerfactory/graphs/mipnerf.py @@ -37,7 +37,7 @@ from nerfactory.renderers.renderers import ( DepthRenderer, RGBRenderer, ) -from nerfactory.utils import colors, visualization, writer +from nerfactory.utils import colors, misc, visualization, writer class MipNerfGraph(Graph): @@ -134,11 +134,12 @@ class MipNerfGraph(Graph): } return outputs - def get_loss_dict(self, outputs, batch): + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): 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) return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): diff --git a/nerfactory/graphs/mipnerf_360.py b/nerfactory/graphs/mipnerf_360.py index 8d699183d..447ebdad6 100644 --- a/nerfactory/graphs/mipnerf_360.py +++ b/nerfactory/graphs/mipnerf_360.py @@ -39,7 +39,7 @@ from nerfactory.renderers.renderers import ( DepthRenderer, RGBRenderer, ) -from nerfactory.utils import colors, visualization, writer +from nerfactory.utils import colors, misc, visualization, writer class MipNerf360Graph(Graph): @@ -143,7 +143,7 @@ class MipNerf360Graph(Graph): } return outputs - def get_loss_dict(self, outputs, batch): + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> 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"]) @@ -153,6 +153,7 @@ class MipNerf360Graph(Graph): "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) return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): diff --git a/nerfactory/graphs/nerfw.py b/nerfactory/graphs/nerfw.py index 8a7d9b7b7..c6f01fbf4 100644 --- a/nerfactory/graphs/nerfw.py +++ b/nerfactory/graphs/nerfw.py @@ -35,7 +35,7 @@ from nerfactory.renderers.renderers import ( RGBRenderer, UncertaintyRenderer, ) -from nerfactory.utils import colors, visualization, writer +from nerfactory.utils import colors, misc, visualization, writer class NerfWGraph(Graph): @@ -180,7 +180,7 @@ class NerfWGraph(Graph): } return outputs - def get_loss_dict(self, outputs, batch): + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): device = outputs["rgb_coarse"].device image = batch["image"].to(device) rgb_coarse = outputs["rgb_coarse"] @@ -198,6 +198,7 @@ class NerfWGraph(Graph): "uncertainty_loss": uncertainty_loss, "density_loss": density_loss, } + loss_dict = misc.scale_dict(loss_dict, loss_coefficients) return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): diff --git a/nerfactory/graphs/semantic_nerf.py b/nerfactory/graphs/semantic_nerf.py index c9759223f..95d9f7d66 100644 --- a/nerfactory/graphs/semantic_nerf.py +++ b/nerfactory/graphs/semantic_nerf.py @@ -36,7 +36,7 @@ from nerfactory.fields.modules.mlp import MLP from nerfactory.fields.nerf_field import NeRFField from nerfactory.graphs.vanilla_nerf import NeRFGraph from nerfactory.renderers.renderers import SemanticRenderer -from nerfactory.utils import writer +from nerfactory.utils import misc, writer class SemanticNerfField(Field): @@ -167,7 +167,7 @@ class SemanticNerfGraph(NeRFGraph): } return outputs - def get_loss_dict(self, outputs, batch): + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients): image = batch["image"] rgb_loss_coarse = self.rgb_loss(image, outputs["rgb_coarse"]) rgb_loss_fine = self.rgb_loss(image, outputs["rgb_fine"]) @@ -179,6 +179,7 @@ class SemanticNerfGraph(NeRFGraph): "rgb_loss_fine": rgb_loss_fine, "semantic_loss_fine": semantic_loss_fine, } + loss_dict = misc.scale_dict(loss_dict, loss_coefficients) return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): diff --git a/nerfactory/graphs/vanilla_nerf.py b/nerfactory/graphs/vanilla_nerf.py index fdf7ba6fc..5604f731a 100644 --- a/nerfactory/graphs/vanilla_nerf.py +++ b/nerfactory/graphs/vanilla_nerf.py @@ -38,7 +38,7 @@ from nerfactory.renderers.renderers import ( DepthRenderer, RGBRenderer, ) -from nerfactory.utils import colors, visualization, writer +from nerfactory.utils import colors, misc, visualization, writer from nerfactory.utils.callbacks import Callback @@ -167,7 +167,8 @@ class NeRFGraph(Graph): } return outputs - def get_loss_dict(self, outputs, batch) -> Dict[str, torch.tensor]: + def get_loss_dict(self, outputs, batch, metrics_dict, loss_coefficients) -> Dict[str, torch.tensor]: + # Scaling metrics by coefficients to create the losses. device = outputs["rgb_coarse"].device image = batch["image"].to(device) @@ -175,6 +176,7 @@ class NeRFGraph(Graph): 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) return loss_dict def log_test_image_outputs(self, image_idx, step, batch, outputs): diff --git a/nerfactory/utils/misc.py b/nerfactory/utils/misc.py index cf26b02e2..69f5a1333 100644 --- a/nerfactory/utils/misc.py +++ b/nerfactory/utils/misc.py @@ -126,3 +126,19 @@ def human_format(num): k = 1000.0 magnitude = int(floor(log(num, k))) return f"{(num / k**magnitude):.2f} {units[magnitude]}" + + +def scale_dict(dictionary: Dict[Any, Any], coefficients: DictConfig) -> Dict[Any, Any]: + """Scale a dictionary in-place given a coefficients dictionary. + + Args: + dictionary: input dict to be scaled. + coefficients: scalar dict config for holding coefficients. + + Returns: + Input dict scaled by coefficients. + """ + for key in dictionary: + if key in coefficients: + dictionary[key] *= coefficients[key] + return dictionary