From 960ac94e06cfc915f08ae00611886413cd1b7ef3 Mon Sep 17 00:00:00 2001
From: Alexander Kristoffersen <alexanator28@icloud.com>
Date: Fri, 22 Jul 2022 18:05:03 -0700
Subject: [PATCH] Loss Metric Refactoring (#181)

* metrics also returned on forward pass, reduction moved to right before backprop
---
 pyrad/engine/trainer.py      | 14 +++++++++-----
 pyrad/graphs/base.py         | 25 +++++++++++++------------
 pyrad/graphs/vanilla_nerf.py |  7 +++++--
 3 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/pyrad/engine/trainer.py b/pyrad/engine/trainer.py
index 20a065634..fdf824788 100644
--- a/pyrad/engine/trainer.py
+++ b/pyrad/engine/trainer.py
@@ -116,11 +116,11 @@ class Trainer:
                     ray_indices, batch = next(iter_dataloader_train)
 
                 with TimeWriter(writer, EventName.ITER_TRAIN_TIME, step=step) as t:
-                    loss_dict = self.train_iteration(ray_indices, batch, step)
+                    loss_metric_dict = self.train_iteration(ray_indices, batch, step)
                 writer.put_scalar(name=EventName.RAYS_PER_SEC, scalar=ray_indices.shape[0] / t.duration, step=step)
 
                 if step != 0 and step % self.config.logging.steps_per_log == 0:
-                    writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_dict, step=step)
+                    writer.put_dict(name="Loss/train-loss_dict", scalar_dict=loss_metric_dict, step=step)
                 if step != 0 and self.config.trainer.steps_per_save and step % self.config.trainer.steps_per_save == 0:
                     self._save_checkpoint(self.config.trainer.model_dir, step)
                 if step % self.config.trainer.steps_per_test == 0:
@@ -188,12 +188,12 @@ class Trainer:
             step: Current training step.
 
         Returns:
-            Dict[str, float]: Dictionary of model losses.
+            Dict[str, float]: Dictionary of model losses and metrics.
         """
         self.optimizers.zero_grad_all()
         with torch.autocast(device_type=ray_indices.device.type, enabled=self.mixed_precision):
-            _, loss_dict = self.graph.forward(ray_indices, batch=batch)
-            loss = loss_dict["aggregated_loss"]
+            _, loss_dict, metrics_dict = self.graph.forward(ray_indices, batch=batch)
+            loss = sum(loss_dict.values())
         self.grad_scaler.scale(loss).backward()
         self.optimizers.optimizer_scaler_step_all(self.grad_scaler)
         self.grad_scaler.update()
@@ -202,6 +202,10 @@ class Trainer:
         if self.graph.callbacks:
             for func_ in self.graph.callbacks:
                 func_.after_step(step)
+
+        # Merging loss and metrics dict into a single output.
+        loss_dict["loss"] = loss
+        loss_dict.update(metrics_dict)
         return loss_dict
 
     @profiler.time_function
diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py
index 5187039cc..5e8d30b7d 100644
--- a/pyrad/graphs/base.py
+++ b/pyrad/graphs/base.py
@@ -163,9 +163,14 @@ class Graph(AbstractGraph):
         masked_intersected_ray_bundle = intersected_ray_bundle[valid_mask]
         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)
-        aggregated_loss_dict = self.get_aggregated_loss_dict(loss_dict)
-        return outputs, aggregated_loss_dict
+
+        # 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]
+        return outputs, loss_dict, metrics_dict
 
     def forward(self, ray_indices: TensorType["num_rays", 3], batch: Union[str, Dict[str, torch.tensor]] = None):
         """Run the forward starting with ray indices."""
@@ -174,17 +179,13 @@ class Graph(AbstractGraph):
 
     @abstractmethod
     def get_loss_dict(self, outputs, batch) -> Dict[str, torch.tensor]:
-        """Computes and returns the losses."""
+        """Computes and returns the losses dict."""
 
-    def get_aggregated_loss_dict(self, loss_dict):
-        """Computes the aggregated loss from the loss_dict and the coefficients specified."""
-        aggregated_loss_dict = {}
-        for loss_name, loss_value in loss_dict.items():
-            assert loss_name in self.loss_coefficients, f"{loss_name} no in self.loss_coefficients"
-            loss_coefficient = self.loss_coefficients[loss_name]
-            aggregated_loss_dict[loss_name] = loss_coefficient * loss_value
-        aggregated_loss_dict["aggregated_loss"] = sum(loss_dict.values())
-        return aggregated_loss_dict
+    def get_metrics_dict(self, outputs, batch) -> Dict[str, torch.tensor]:
+        """Compute and obtain metrics and coefficients."""
+        # pylint: disable=unused-argument
+        # pylint: disable=no-self-use
+        return {}
 
     @torch.no_grad()
     def get_outputs_for_camera_ray_bundle(self, camera_ray_bundle: RayBundle):
diff --git a/pyrad/graphs/vanilla_nerf.py b/pyrad/graphs/vanilla_nerf.py
index 54bf56bb2..ad28e7e7a 100644
--- a/pyrad/graphs/vanilla_nerf.py
+++ b/pyrad/graphs/vanilla_nerf.py
@@ -140,10 +140,13 @@ class NeRFGraph(Graph):
         }
         return outputs
 
-    def get_loss_dict(self, outputs, batch):
-        image = batch["image"]
+    def get_loss_dict(self, outputs, batch) -> Dict[str, torch.tensor]:
+        device = outputs["rgb_coarse"].device
+        image = batch["image"].to(device)
+
         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}
         return loss_dict