Issue with second derivatives using tiny-cuda-nn hash encoding

Question

Issue with second derivatives using tiny-cuda-nn hash encoding

jfrausto7 opened this issue 9 months ago · 2 comments

Hi there! I'm having trouble using tiny-cuda-nn's hash encoding when computing a Hessian matrix. As mentioned in other open & closed issues, second derivatives are not supported by tiny-cuda-nn's neural networks. I'm building on top of a repository that uses tiny-cuda-nn (torch-ngp), and I'm trying to compose the hash encoding (tcnn.Encoding) in a network with a simple MLP as shown below (link to full code here):

class MLP(nn.Module):
    def __init__(self, dim_in, dim_out, dim_hidden, num_layers, bias=True):
        super().__init__()
        self.dim_in = dim_in
        self.dim_out = dim_out
        self.dim_hidden = dim_hidden
        self.num_layers = num_layers

        net = []
        for l in range(num_layers):
            net.append(nn.Linear(self.dim_in if l == 0 else self.dim_hidden, self.dim_out if l == num_layers - 1 else self.dim_hidden, bias=bias))

        self.net = nn.ModuleList(net)
    
    def forward(self, x):
        for l in range(self.num_layers):
            x = self.net[l](x)
            if l != self.num_layers - 1:
                x = F.relu(x, inplace=True)
        return x

class NeRFNetwork(NeRFRenderer):
    def __init__(self,
                 encoding="HashGrid",
                 encoding_dir="SphericalHarmonics",
                 num_layers=2,
                 hidden_dim=64,
                 geo_feat_dim=15,
                 num_layers_color=3,
                 hidden_dim_color=64,
                 bound=1,
                 **kwargs
                 ):
        super().__init__(bound, **kwargs)

        # sigma network
        self.num_layers = num_layers
        self.hidden_dim = hidden_dim
        self.geo_feat_dim = geo_feat_dim

        per_level_scale = np.exp2(np.log2(2048 * bound / 16) / (16 - 1))

        self.encoder = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "HashGrid",
                "n_levels": 16,
                "n_features_per_level": 2,
                "log2_hashmap_size": 19,
                "base_resolution": 16,
                "per_level_scale": per_level_scale,
            },
        )

        self.sigma_net = MLP(dim_in=32, dim_out=1 + self.geo_feat_dim, dim_hidden=self.hidden_dim, num_layers=self.num_layers)

        # color network
        self.num_layers_color = num_layers_color        
        self.hidden_dim_color = hidden_dim_color

        self.encoder_dir = tcnn.Encoding(
            n_input_dims=3,
            encoding_config={
                "otype": "SphericalHarmonics",
                "degree": 4,
            },
        )

        self.in_dim_color = self.encoder_dir.n_output_dims + self.geo_feat_dim

        self.color_net = MLP(dim_in=self.in_dim_color, dim_out=3, dim_hidden=self.hidden_dim_color, num_layers=self.num_layers_color)

    
    def forward(self, x, d):
        # x: [N, 3], in [-bound, bound]
        # d: [N, 3], nomalized in [-1, 1]


        # sigma
        x = (x + self.bound) / (2 * self.bound) # to [0, 1]
        x = self.encoder(x)
        h = self.sigma_net(x)

        #sigma = F.relu(h[..., 0])
        sigma = trunc_exp(h[..., 0])
        geo_feat = h[..., 1:]

        # color
        d = (d + 1) / 2 # tcnn SH encoding requires inputs to be in [0, 1]
        d = self.encoder_dir(d)

        #p = torch.zeros_like(geo_feat[..., :1]) # manual input padding
        h = torch.cat([d, geo_feat], dim=-1)
        h = self.color_net(h)
        
        # sigmoid activation for rgb
        color = torch.sigmoid(h)

        return sigma, color

Even after replacing all instances of `tcnn.Network' with an MLP and rebuilding/training, I still get the following error:

Traceback (most recent call last): File "/home/jfrausto/nerf-navigation/simulate.py", line 350, in <module> simulate(planner_cfg, agent_cfg, filter_cfg, extra_cfg, density_fn, render_fn, get_rays_fn) File "/home/jfrausto/nerf-navigation/simulate.py", line 88, in simulate state_est = filter.estimate_state(gt_img, true_pose, action) File "/home/jfrausto/nerf-navigation/nav/estimator_helpers.py", line 384, in estimate_state hess = torch.autograd.functional.hessian(lambda x: self.measurement_fn(x, self.xt.clone().detach(), sig_prop, self.target, self.batch), xt.clone().detach()) File "/home/jfrausto/.local/lib/python3.10/site-packages/torch/autograd/functional.py", line 812, in hessian res = jacobian(jac_func, inputs, create_graph=create_graph, strict=strict, vectorize=vectorize, File "/home/jfrausto/.local/lib/python3.10/site-packages/torch/autograd/functional.py", line 673, in jacobian vj = _autograd_grad((out.reshape(-1)[j],), inputs, File "/home/jfrausto/.local/lib/python3.10/site-packages/torch/autograd/functional.py", line 159, in _autograd_grad return torch.autograd.grad(new_outputs, inputs, new_grad_outputs, allow_unused=True, File "/home/jfrausto/.local/lib/python3.10/site-packages/torch/autograd/__init__.py", line 300, in grad return Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass File "/home/jfrausto/.local/lib/python3.10/site-packages/torch/autograd/function.py", line 267, in apply return user_fn(self, *args) File "/home/jfrausto/.local/lib/python3.10/site-packages/tinycudann/modules.py", line 145, in backward doutput_grad, params_grad, input_grad = ctx.ctx_fwd.native_tcnn_module.bwd_bwd_input( RuntimeError: DifferentiableObject::backward_backward_input_impl: not implemented error

I'm not sure what I'm doing wrong here, but this seems to still be an issue with tiny-cuda-nn. Would really appreciate any advice. Thanks!

Answer 1 · 2024-04-21T07:41:09.000Z