Repository: nerfstudio-project/nerfstudio

Issue:
instant ngp tcnn error with visualization naans?
![Screenshot from 2022-06-02 13-20-05](https://user-images.githubusercontent.com/14854811/171731758-7c472092-f9cc-4dd8-aa1c-e352c85c524e.png)



Discussion:
Is it only tcnn that causes nans?
yeah it seems like it...
 i wonder if this is half precision related?
I ran into this today and stepped around in PDB a bit; some observations ---

- We're getting a float16 overflow in a forward pass right before the softplus here:
https://github.com/plenoptix/pyrad/blob/5a1c8a1c86bd6ad31f437c19bac9dda4f0d58c36/pyrad/fields/instant_ngp_field.py#L115-L118
- A few position encoding parameters drift close to the 1e4-ish range, which cause an overflow inside the MLP. Results in a NaN loss, usually after somewhere between 50k and 80k steps.

- `torch-ngp` [uses a truncated exponential](https://github.com/ashawkey/torch-ngp/blob/f53441ef094becf03ef0363522087a7733670a6c/nerf/network_tcnn.py#L91-L96) instead of a softplus which I subbed in and helps a lot; intuitively should avoid overflow by enabling really high post-activation (float32) density outputs from smaller internal (float16) parameters. I have an old commit of the original NGP repo that appears to have gotten away with a ReLU, but from a glance the newer version seems to also use a similar exponential ([1](https://github.com/NVlabs/instant-ngp/blob/1a99843f278d88f9c4da414bc3df542539ba4370/src/testbed_nerf.cu#L246), [2](https://github.com/NVlabs/instant-ngp/blob/1a99843f278d88f9c4da414bc3df542539ba4370/include/neural-graphics-primitives/testbed.h#L627))? 
- Test PSNRs look more stable after the change, although the peak is slightly worse; softplus in orange and exp in blue:
  ![image](https://user-images.githubusercontent.com/6992947/177864700-6089a1f7-db8d-45ce-9d38-5b065b551c74.png)

- Still eventually get a NaN, but it's now from an exploding gradient (loss is not NaN but some gradient overflows) and at around 140k steps.
Thanks for digging deeper into this! Exp activation is good find, seem like it is worth changing. Mind if I reassign this issue to you?
Yep can make a PR!

Resolving patch:
diff --git a/pyrad/fields/instant_ngp_field.py b/pyrad/fields/instant_ngp_field.py
--- a/pyrad/fields/instant_ngp_field.py
+++ b/pyrad/fields/instant_ngp_field.py
@@ -20,7 +20,6 @@
 from typing import Tuple
 
 import torch
-import torch.nn.functional as F
 from torch.nn.parameter import Parameter
 
 from pyrad.fields.modules.encoding import Encoding, HashEncoding, SHEncoding
@@ -28,6 +27,7 @@
 from pyrad.fields.base import Field
 from pyrad.fields.nerf_field import NeRFField
 from pyrad.cameras.rays import RaySamples
+from pyrad.utils.activations import trunc_exp
 
 try:
     import tinycudann as tcnn
@@ -115,7 +115,12 @@ def get_density(self, ray_samples: RaySamples):
         x = self.position_encoding(positions_flat)
         h = self.mlp_base(x).view(*ray_samples.frustums.get_positions().shape[:-1], -1).to(dtype)
         density_before_activation, base_mlp_out = torch.split(h, [1, self.geo_feat_dim], dim=-1)
-        density = F.softplus(density_before_activation)
+
+        # Rectifying the density with an exponential is much more stable than a ReLU or
+        # softplus, because it enables high post-activation (float32) density outputs
+        # from smaller internal (float16) parameters.
+        assert density_before_activation.dtype is torch.float32
+        density = trunc_exp(density_before_activation)
         return density, base_mlp_out
 
     def get_outputs(self, ray_samples: RaySamples, density_embedding=None):
diff --git a/pyrad/graphs/base.py b/pyrad/graphs/base.py
--- a/pyrad/graphs/base.py
+++ b/pyrad/graphs/base.py
@@ -18,7 +18,6 @@
 from abc import abstractmethod
 from collections import defaultdict
 import time
-import traceback
 from typing import Any, Dict, List, Union
 
 import torch
diff --git a/pyrad/utils/activations.py b/pyrad/utils/activations.py
new file mode 100644
--- /dev/null
+++ b/pyrad/utils/activations.py
@@ -0,0 +1,42 @@
+# 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.
+
+"""
+Special activation functions.
+"""
+
+import torch
+from torch.autograd import Function
+from torch.cuda.amp import custom_bwd, custom_fwd
+
+
+class _TruncExp(Function):  # pylint: disable=abstract-method
+    # Implementation from torch-ngp:
+    # https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
+    @staticmethod
+    @custom_fwd(cast_inputs=torch.float32)
+    def forward(ctx, x):  # pylint: disable=arguments-differ
+        ctx.save_for_backward(x)
+        return torch.exp(x)
+
+    @staticmethod
+    @custom_bwd
+    def backward(ctx, g):  # pylint: disable=arguments-differ
+        x = ctx.saved_tensors[0]
+        return g * torch.exp(x.clamp(-15, 15))
+
+
+trunc_exp = _TruncExp.apply
+"""Same as torch.exp, but with the backward pass clipped to prevent vanishing/exploding
+gradients."""


Test patch: