Cannot backprop gradients from PyTorch model loss to Mitsuba3

[matthewdhull]

I am making use of drjit.wrap_ad decorators with the intent of taking the x-entropy loss of a PyTorch Image Classifier and backprop'ing the gradient of the loss through Mitsuba to get the grad of the loss w.r.t. to the vertex colors of a mesh. I cannot get gradients for a scene parameter with my current approach.

I have a mesh in my scene with a differentiable sphere_color.vertex_color param

SceneParameters[
  --------------------------------------------------------------------------
  Name                   Flags    Type            Parent
  --------------------------------------------------------------------------
  sphere.vertex_color    ∂        Float           PLYMesh
]

I have the following wrapped funcs to classify the rendered image and calculate the loss:

@dr.wrap_ad(source='drjit', target='torch')
def torch_criterion(logits):
    target = ch.tensor([1]).cuda()
    criterion = nn.CrossEntropyLoss()
    loss = criterion(logits, target)
    return loss


@dr.wrap_ad(source='drjit', target='torch')
def torch_model(model, x):
    t = ch.unsqueeze(x.permute((2,1,0)),0)
    logits = model(t)
    return logits

In my custom optimization loop, I render the image, pass it to the model, obtain classification logits and loss and then call dr.backward(loss) but the gradient is always 0.0 for each vertex. Minimal example of my loop:

def optim(scene, k, iters ):
    params = mi.traverse(scene)
    orig_tex = dr.cuda.ad.Float(params[k]) # save original tex colors
    dr.enable_grad(params[k])
    params.update()
    for it in range(iters):
        print(f'iter {it}')
        loss = mi.Float(0.0)
        img = mi.render(scene)
        logits = torch_model(model, img)
        loss = torch_criterion(logits).array
        dr.enable_grad(loss)
        dr.backward(loss)
        grad = dr.grad(params[k])  # grad is always [0.0, 0.0, .... 0.0] for each vertex
        . . . 

optim(scene, k='sphere.vertex_color', iters=50)

Thank you for any assistance you can provide in resolving the gradient issue.

[Speierers]

Could you make sure you get some gradients using a simple loss function implemented in Dr.Jit? E.g. loss = dr.mean(dr.sqr(img)).

Also, I would recommend you try combining both AD wrappers into one, maybe this is causing some issues here?

[matthewdhull]

Unfortunately, I am not getting gradients with loss = dr.mean(dr.sqr(img)) either.

If it's helpful, the initial rendering of the sphere contains an initialization of random colors for each vertex:

[Speierers]

Could you try passing params to the mi.render call in the optimization loop like this: img = mi.render(scene, params)?

[matthewdhull]

@Speierers, I was able to get the gradients using img = mi.render(scene, params) or the simple loss calculation loss = dr.mean(dr.sqr(img))

With that success, I followed your suggestion in combining the the AD wrapped functions into one method:

@dr.wrap_ad(source='drjit', target='torch')
def torch_criterion(model, x):
    x = ch.tensor(x.detach().numpy())
    t = ch.unsqueeze(x.permute((2,1,0)),0)
    logits = model(t)
    target = ch.tensor([1]).cuda() 
    criterion = nn.CrossEntropyLoss()
    loss = criterion(logits, target)
    return loss

the loss variable is torch.float32. When the decorator is invoked, I now get:

Critical failure in Dr.Jit AD backend: referenced an unknown variable a4!

When I step through the code, the error happens here:

Snippet of code at line 5654 in router.py

        Type.add_edge_(
            diff_vars_in[0]  if len(diff_vars_in)  == 1 else tmp_in.index_ad,
            diff_vars_out[0] if len(diff_vars_out) == 1 else tmp_out.index_ad,
            inst
        )

[Speierers]

Could you tell whether this concerns an input or output variable?

Also is it necessary that you pass the model to the wrapper fonction?

If you could provide a minimal reproducer for this issue I could take a closer look on my side.