diff --git a/pyrad/fields/instant_ngp_field.py b/pyrad/fields/instant_ngp_field.py index 12bc0bfc9..c62fa3de9 100644 --- a/pyrad/fields/instant_ngp_field.py +++ b/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): diff --git a/pyrad/utils/activations.py b/pyrad/utils/activations.py new file mode 100644 index 000000000..097f83dca --- /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."""