The backward error of SphericalHarmonics encoding

[lzhnb]

Hi Thomas,

Thanks for your amazing work.

I met some problems when I try to backward the output of the SphericalHarmonics module. Here is the code:

import torch
import torch.nn.functional as F
import tinycudann as tcnn

sh_enc_degree = 4

sh_encoder = tcnn.Encoding(
    n_input_dims=3,
    encoding_config={"otype": "SphericalHarmonics", "degree": sh_enc_degree},
    dtype=torch.float32,
)

inputs = F.normalize(torch.randn(100, 3), dim=-1).cuda()
inputs.requires_grad_()
outputs = sh_encoder(inputs)

outputs.backward(torch.zeros_like(outputs))

And here is the error in console:

Traceback (most recent call last):
  File "test_sh_encoder.py", line 17, in <module>
    outputs.backward(torch.zeros_like(outputs))
  File "xxx/site-packages/torch/tensor.py", line 245, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph, inputs=inputs)
  File "xxx/site-packages/torch/autograd/__init__.py", line 145, in backward
    Variable._execution_engine.run_backward(
  File "xxx/site-packages/torch/autograd/function.py", line 89, in apply
    return self._forward_cls.backward(self, *args)  # type: ignore
  File "xxx/site-packages/torch/autograd/function.py", line 210, in wrapper
    outputs = fn(ctx, *args)
  File "xxx/tiny-cuda-nn/bindings/torch/tinycudann/modules.py", line 52, in backward
    input_grad, weight_grad = ctx.native_tcnn_module.bwd(ctx.native_ctx, input, params, output, scaled_grad)
RuntimeError: Encoding: forward(prepare_input_gradients) must be called before backward(dL_dinput)

I found that you have realized the kernel_sh_backward in spherical_harmonics.h. It maybe some problems in pytorch autograd wrapper?

Otherwise, the backward of HashGrid is bug-free, thanks :)

[Tom94]

Hi Zhihao,

your guess about the bindings being at fault is correct! The changes behind https://github.com/NVlabs/tiny-cuda-nn/pull/62 should also have fixed this problem. Could you pull master and let me know?

Cheers!

[lzhnb]

It works👍! Thanks!