From 02f449b661e1578958d49b60628e82978de09e60 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jon=C3=A1=C5=A1=20Kulh=C3=A1nek?= <jonas.kulhanek@live.com>
Date: Fri, 12 May 2023 07:23:41 +0200
Subject: [PATCH] Implement gradient scaling (#1901)

* Implement gradient scaling

* Make grad scaling more general

* Fix pylint errors

* Fix lint error + change gradient scaling

* Fix issues with grad scaling

* Fix linting errors
---
 nerfstudio/model_components/losses.py | 39 ++++++++++++++++++++++++++-
 nerfstudio/models/nerfacto.py         |  9 ++++++-
 2 files changed, 46 insertions(+), 2 deletions(-)

diff --git a/nerfstudio/model_components/losses.py b/nerfstudio/model_components/losses.py
index 43fbbd271..d46193341 100644
--- a/nerfstudio/model_components/losses.py
+++ b/nerfstudio/model_components/losses.py
@@ -16,7 +16,7 @@
 Collection of Losses.
 """
 from enum import Enum
-from typing import Literal
+from typing import Dict, Literal
 
 import torch
 from torch import nn
@@ -502,3 +502,40 @@ def tv_loss(grids: TensorType["grids", "feature_dim", "row", "column"]) -> Tenso
     h_tv = torch.pow((grids[:, :, 1:, :] - grids[:, :, :-1, :]), 2).sum()
     w_tv = torch.pow((grids[:, :, :, 1:] - grids[:, :, :, :-1]), 2).sum()
     return 2 * (h_tv / h_tv_count + w_tv / w_tv_count) / number_of_grids
+
+
+class _GradientScaler(torch.autograd.Function):  # typing: ignore, pylint: disable=abstract-method
+    """
+    Scale gradients by a constant factor.
+    """
+
+    @staticmethod
+    def forward(ctx, value, scaling):  # pylint: disable=arguments-differ
+        ctx.save_for_backward(scaling)
+        return value, scaling
+
+    @staticmethod
+    def backward(ctx, output_grad, grad_scaling):  # pylint: disable=arguments-differ
+        (scaling,) = ctx.saved_tensors
+        return output_grad * scaling, grad_scaling
+
+
+def scale_gradients_by_distance_squared(
+    field_outputs: Dict[str, torch.Tensor], ray_samples: RaySamples
+) -> Dict[str, torch.Tensor]:
+    """
+    Scale gradients by the ray distance to the pixel
+    as suggested in `Radiance Field Gradient Scaling for Unbiased Near-Camera Training` paper
+
+    Note: The scaling is applied on the interval of [0, 1] along the ray!
+
+    Example:
+        GradientLoss should be called right after obtaining the densities and colors from the field. ::
+            >>> field_outputs = scale_gradient_by_distance_squared(field_outputs, ray_samples)
+    """
+    out = {}
+    ray_dist = (ray_samples.frustums.starts + ray_samples.frustums.ends) / 2
+    scaling = torch.square(ray_dist).clamp(0, 1)
+    for key, value in field_outputs.items():
+        out[key], _ = _GradientScaler.apply(value, scaling)
+    return out
diff --git a/nerfstudio/models/nerfacto.py b/nerfstudio/models/nerfacto.py
index 0a7583ebe..858858c19 100644
--- a/nerfstudio/models/nerfacto.py
+++ b/nerfstudio/models/nerfacto.py
@@ -28,7 +28,7 @@ from torchmetrics import PeakSignalNoiseRatio
 from torchmetrics.functional import structural_similarity_index_measure
 from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
 
-from nerfstudio.cameras.rays import RayBundle
+from nerfstudio.cameras.rays import RayBundle, RaySamples
 from nerfstudio.engine.callbacks import (
     TrainingCallback,
     TrainingCallbackAttributes,
@@ -44,6 +44,7 @@ from nerfstudio.model_components.losses import (
     interlevel_loss,
     orientation_loss,
     pred_normal_loss,
+    scale_gradients_by_distance_squared,
 )
 from nerfstudio.model_components.ray_samplers import (
     ProposalNetworkSampler,
@@ -127,6 +128,8 @@ class NerfactoModelConfig(ModelConfig):
     """Whether to predict normals or not."""
     disable_scene_contraction: bool = False
     """Whether to disable scene contraction or not."""
+    use_gradient_scaling: bool = False
+    """Use gradient scaler where the gradients are lower for points closer to the camera."""
 
 
 class NerfactoModel(Model):
@@ -262,8 +265,12 @@ class NerfactoModel(Model):
         return callbacks
 
     def get_outputs(self, ray_bundle: RayBundle):
+        ray_samples: RaySamples
         ray_samples, weights_list, ray_samples_list = self.proposal_sampler(ray_bundle, density_fns=self.density_fns)
         field_outputs = self.field(ray_samples, compute_normals=self.config.predict_normals)
+        if self.config.use_gradient_scaling:
+            field_outputs = scale_gradients_by_distance_squared(field_outputs, ray_samples)
+
         weights = ray_samples.get_weights(field_outputs[FieldHeadNames.DENSITY])
         weights_list.append(weights)
         ray_samples_list.append(ray_samples)