Truncated exponential for fixing NaN in instant-ngp (#104)

* replace ngp softplus with exponential, fixes #11

* resolve pylint errors

pyrad/fields/instant_ngp_field.py

@@ -20,7 +20,6 @@ Instant-NGP field implementations using tiny-cuda-nn, torch, ....
 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.modules.field_heads import FieldHeadNames
 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 @@ class TCNNInstantNGPField(Field):
         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):

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