Add interval loss (#92)

Implement interval loss that was proposed in mipnerf-360.

configs/graph_mipnerf_360.yaml

@@ -27,4 +27,7 @@ graph:
   collider_config:
     _target_: pyrad.graphs.modules.scene_colliders.NearFarCollider
     near_plane: 0.5
-    far_plane: 20.0
+    far_plane: 20.0
+  loss_coefficients:
+    ray_loss_coarse: 1.0
+    ray_loss_fine: 1.0

docs/reference/graphs/modules/index.rst

@@ -7,4 +7,5 @@ TODO: High level description of graph modules and how they are used.
 
 .. toctree::
 
-   ray_sampler
+   ray_sampler
+   ray_losses

docs/reference/graphs/modules/ray_losses.rst

@@ -0,0 +1,8 @@
+.. _ray_losses:
+
+Ray Losses
+===================
+
+.. automodule:: pyrad.graphs.modules.ray_losses
+   :members:
+   :show-inheritance:

pyrad/graphs/mipnerf_360.py

@@ -30,6 +30,7 @@ from pyrad.fields.modules.field_heads import FieldHeadNames
 from pyrad.fields.modules.spatial_distortions import SceneContraction
 from pyrad.fields.nerf_field import NeRFField
 from pyrad.graphs.base import Graph
+from pyrad.graphs.modules.ray_losses import interval_loss
 from pyrad.optimizers.loss import MSELoss
 from pyrad.graphs.modules.ray_sampler import PDFSampler, UniformSampler
 from pyrad.renderers.renderers import AccumulationRenderer, DepthRenderer, RGBRenderer
@@ -111,6 +112,7 @@ class MipNerf360Graph(Graph):
         )
         accumulation_coarse = self.renderer_accumulation(weights_coarse)
         depth_coarse = self.renderer_depth(weights_coarse, ray_samples_uniform)
+        ray_loss_coarse = interval_loss(ray_samples_uniform, field_outputs_coarse[FieldHeadNames.DENSITY])
 
         # pdf sampling
         ray_samples_pdf = self.sampler_pdf(ray_bundle, ray_samples_uniform, weights_coarse)
@@ -124,6 +126,7 @@ class MipNerf360Graph(Graph):
         )
         accumulation_fine = self.renderer_accumulation(weights_fine)
         depth_fine = self.renderer_depth(weights_fine, ray_samples_pdf)
+        ray_loss_fine = interval_loss(ray_samples_pdf, field_outputs_fine[FieldHeadNames.DENSITY])
 
         outputs = {
             "rgb_coarse": rgb_coarse,
@@ -132,6 +135,8 @@ class MipNerf360Graph(Graph):
             "accumulation_fine": accumulation_fine,
             "depth_coarse": depth_coarse,
             "depth_fine": depth_fine,
+            "ray_loss_coarse": ray_loss_coarse,
+            "ray_loss_fine": ray_loss_fine,
         }
         return outputs
 
@@ -139,7 +144,12 @@ class MipNerf360Graph(Graph):
         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 = {
+            "rgb_loss_coarse": rgb_loss_coarse,
+            "rgb_loss_fine": rgb_loss_fine,
+            "ray_loss_coarse": torch.mean(outputs["ray_loss_coarse"]),
+            "ray_loss_fine": torch.mean(outputs["ray_loss_fine"]),
+        }
         return loss_dict
 
     def log_test_image_outputs(self, image_idx, step, batch, outputs):
@@ -183,6 +193,12 @@ class MipNerf360Graph(Graph):
         writer.put_scalar(name=f"psnr/val_{image_idx}-fine", scalar=float(fine_psnr), step=step)
         writer.put_scalar(name=f"ssim/val_{image_idx}", scalar=float(fine_ssim), step=step)
         writer.put_scalar(name=f"lpips/val_{image_idx}", scalar=float(fine_lpips), step=step)
+        writer.put_scalar(
+            name=f"ray_loss_coarse/val_{image_idx}", scalar=float(torch.mean(outputs["ray_loss_coarse"])), step=step
+        )
+        writer.put_scalar(
+            name=f"ray_loss_fine/val_{image_idx}", scalar=float(torch.mean(outputs["ray_loss_fine"])), step=step
+        )
 
         writer.put_scalar(name=writer.EventName.CURR_TEST_PSNR, scalar=float(fine_psnr), step=step)
 

pyrad/graphs/modules/ray_losses.py

@@ -0,0 +1,53 @@
+# Copyright 2022 The Plenoptix Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Collection of Ray Losses"""
+
+import torch
+from torchtyping import TensorType
+
+from pyrad.cameras.rays import RaySamples
+
+
+def interval_loss(ray_samples: RaySamples, densities: TensorType[..., "num_samples", 1]) -> TensorType[..., 1]:
+    """Interval loss origially proposed in MipNeRF-360.
+
+    .. math::
+
+        \\mathcal{L}(\\mathbf{s}, \\mathbf{w}) =\\iint\\limits_{-\\infty}^{\\,\\,\\,\\infty}
+        \\mathbf{w}_\\mathbf{s}(u)\\mathbf{w}_\\mathbf{s}(v)|u - v|\\,d_{u}\\,d_{v}
+
+    where :math:`\\mathbf{w}_\\mathbf{s}(u)=\\sum_i w_i \\mathbb{1}_{[\\mathbf{s}_i, \\mathbf{s}_{i+1})}(u)`
+    is the weight at location :math:`u` between bin locations :math:`s_i` and :math:`s_{i+1}`.
+
+    Args:
+        ray_samples (RaySamples): Ray samples to compute loss over
+        densities (TensorType[..., "num_samples", 1]): Predicted sample densities
+
+    Returns:
+        TensorType[..., 1]: Interval Loss.
+    """
+
+    # Compute the weight at each sample location
+    weights = ray_samples.get_weights(densities)
+
+    midpoints = (ray_samples.bin_starts + ray_samples.bin_ends) / 2.0  # (..., num_samples, 1)
+
+    loss = (
+        weights * weights[..., None, :, 0] * torch.abs(midpoints - midpoints[..., None, :, 0])
+    )  # (..., num_samples, num_samples)
+    loss = torch.sum(loss, dim=(-1, -2))[..., None]  # (..., num_samples)
+    loss = loss + 1 / 3.0 * torch.sum(weights**2 * (ray_samples.bin_ends - ray_samples.bin_starts), dim=-2)
+
+    return loss