Neural SDF

[zichenwang01]

Hi!

I wanted to code up a neural SDF differentiable renderer but am not too sure how to do this with Mitsuba. Specifically, I want to instantiate the neural SDF in PyTorch but code the differentiable renderer part in Mitsuba, and in the end my goal is to optimize the neural SDF.

I took a detailed read on this and came up with the following: I tried wrapping the torch neural sdf class with a Mitsuba class, so that in my differentiable render I can call the sdf methods as if everything is in Mitsuba

class NeuralSDFModule(nn.Module):
    def __init__(self, input_dim=3, hidden_dim=128, output_dim=1, num_layers=4):
        super().__init__()
        layers = []
        for i in range(num_layers):
            layers.append(nn.Linear(input_dim if i == 0 else hidden_dim, hidden_dim))
            layers.append(nn.ReLU())
        layers.append(nn.Linear(hidden_dim, output_dim))
        self.model = nn.Sequential(*layers)

    def forward(self, x):
        return self.model(x)
    
class NeuralSDF(mi.Object):
    def __init__(self):
        super().__init__()
        self.module = NeuralSDFModule().to('cuda')
        self.bbox = mi.BoundingBox3f(mi.Point3f(0, 0, 0), mi.Point3f(1, 1, 1))  
        
    @dr.wrap(source='drjit', target='torch')
    def at(self, p:Point3f) -> Float:
        """Evaluate the SDF value at given positions."""
        p = p.permute(1, 0)
        return self.module(p)

but then I realized I needed a torch optimizer to optimize the torch module, which means that my loss also needs to be a torch tensor. So I did more tricks to try to come around this--

sdf = NeuralSDF()

print("----- finish load sdf -----")

import torch
optimizer = torch.optim.Adam(sdf.module.parameters(), lr=1e-3)

print("----- finish load optimizer -----")

@dr.wrap(source='torch', target='drjit')
def compute_loss():
    p = Point3f(0.5, 0.5, 0.5)
    loss = sdf.at(p) # mimicking the differentiable renderer in Mitsuba
    return loss

loss = compute_loss()
print("loss (torch): ", loss)

print("----- finish loss -----")

optimizer.zero_grad()
loss.backward()
optimizer.step()

print("----- finish step -----")

But anyhow right now I am getting the following error message--

----- finish load sdf -----
----- finish load optimizer -----
torch.Size([1, 3])
loss (torch):  tensor([[0.0197]], device='cuda:0')
----- finish loss -----
Traceback (most recent call last):
  File "/nfs/turbo/jjparkcv-turbo-large/zichen/differentiable-rendering/tmp.py", line 26, in <module>
    loss.backward()
  File "/nfs/turbo/jjparkcv-turbo-large/zichen/anaconda3/envs/dr2/lib/python3.10/site-packages/torch/_tensor.py", line 581, in backward
    torch.autograd.backward(
  File "/nfs/turbo/jjparkcv-turbo-large/zichen/anaconda3/envs/dr2/lib/python3.10/site-packages/torch/autograd/__init__.py", line 347, in backward
    _engine_run_backward(
  File "/nfs/turbo/jjparkcv-turbo-large/zichen/anaconda3/envs/dr2/lib/python3.10/site-packages/torch/autograd/graph.py", line 825, in _engine_run_backward
    return Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass
RuntimeError: element 0 of tensors does not require grad and does not have a grad_fn

This seems a bit complicated and I am not sure what is the correct to do this in Mitsuba. Has anyone in the Mitsuba team tried similar things before? Any suggestions on how to make this work would be greatly appreciated!

Happy holidays,
Zichen

[njroussel]

I'm not familiar enough with torch to give you pointers to what could be going wrong here just based off of that error message.

One thing to note though, depending on how you wrote the rendering code itself, you will need be using the evaluated mode of Dr.Jit. This wasn't necessary in the tutorial you linked, because the neural parts happened after the render process, and not in middle of it as is the case here.