Double backwards support for tcnn.Network
adamdai opened this issue · comments
Hello,
I am using a tcnn Hash Encoding + Cutlass MLP to train a neural field. I would like to access second derivatives to implement a smoothing loss during training, as well as obtain gradients for performing trajectory optimization over the neural field.
I noticed that double backwards support has been implemented for tcnn.Encoding here. However, the example uses a torch MLP which is much slower. When trying to obtain 2nd derivatives with a tcnn.Network MLP instead, a backward_backward_input_impl: not implemented error
is thrown.
Here is a minimal example for testing:
import torch
import tinycudann as tcnn
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class NeuralField(torch.nn.Module):
def __init__(self):
super().__init__()
self.encoding = tcnn.Encoding(
n_input_dims=2,
encoding_config={
"otype": "HashGrid",
"n_levels": 8,
"n_features_per_level": 8,
"log2_hashmap_size": 19,
"base_resolution": 16,
"per_level_scale": 1.2599210739135742,
"interpolation": "Smoothstep"
},
)
tot_out_dims_2d = self.encoding.n_output_dims
self.mlp = tcnn.Network(
n_input_dims=tot_out_dims_2d,
n_output_dims=1,
network_config={
"otype": "CutlassMLP",
"activation": "ReLU",
"output_activation": "None",
"n_neurons": 256,
"n_hidden_layers": 3,
},
)
def forward(self, x):
x = self.encoding(x)
x = self.mlp(x)
return x
if __name__ == "__main__":
nf = NeuralField()
nf.to(device)
x = torch.rand(10, 2, requires_grad=True).to(device)
y = nf(x)
# 2 methods of obtaining 2nd derivative - both give the same error
method = 0
if method == 0:
grad = torch.autograd.grad(y.sum(), x, create_graph=True)[0]
grad_2 = torch.autograd.grad(grad.sum(), x)[0]
else:
grad = torch.autograd.grad(y, x, torch.ones_like(y, device=x.device),
create_graph=True, retain_graph=True, only_inputs=True)[0]
grad_2 = torch.autograd.grad(grad, x, torch.ones_like(grad, device=x.device),
create_graph=False, retain_graph=False, only_inputs=True)[0]
The full error trace is:
Traceback (most recent call last):
File "tcnn_double_backward.py", line 54, in <module>
grad_2 = torch.autograd.grad(grad.sum(), x)[0]
File "/home/navlab/anaconda3/envs/nemo/lib/python3.8/site-packages/torch/autograd/__init__.py", line 394, in grad
result = Variable._execution_engine.run_backward( # Calls into the C++ engine to run the backward pass
File "/home/navlab/anaconda3/envs/nemo/lib/python3.8/site-packages/torch/autograd/function.py", line 288, in apply
return user_fn(self, *args)
File "/home/navlab/anaconda3/envs/nemo/lib/python3.8/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
Are there any plans to implement double backward for tcnn.Network? It would be greatly helpful for my project.