Updating metrics dict code and logic for all the graphs (#198)

* illustrating logic of metrics dict with vanilla nerf

* updated all graph impls

* clean up.

* cleanup.

* format.

Co-authored-by: Alexander Kristoffersen <akristoffersen@berkeley.edu>

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."""

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

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):

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):

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):

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):

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):

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