How to optimize brdf textures?

[osylum]

Hi,

I have a custom BSDF with textures (BitmapTexture) and try to optimize their values in a optimization loop. Checking the scene parameters, I see that they are flagged as differentiable as well as the corresponding data of the BitmapTexture. But I get the following error at the dr.backward call:

TypeError: backward_from(): the argument does not depend on the input variable(s) being differentiated. Raising an exception since this is usually indicative of a bug (for example, you may have forgotten to call dr.enable_grad(..)). If this is expected behavior, skip the call to backward_from(..) if ek.grad_enabled(..) returns False.
drjit-autodiff: variable leak detected (4 variables remain in use)!

• I get the texture parameter as: params = mi.traverse(scene_params[key])
• Note that my brdf renders fine and I can change the texture fine when assigning a new mi.TensorXf to params['data']
• I keep gradient in the param as: params.keep('data')
• I enabled gradient as: dr.enable_grad(params['data'])
• I assign the parameter to optimize as: opt[key] = params['data']
• I render and backpropagate as:
  • image = mi.render(scene, scene_params)
  • loss = dr.mean(dr.abs(image[:,:,0:3] - images_true[i]))
  • dr.backward_from(loss) # tried also with ADFlag.ClearNone

The scene parameter is:

scene_params[key_rho_d]=BitmapTexture[
  name = "rho_d.exr",
  resolution = "[256, 256]",
  raw = 1,
  mean = [0.5],
  transform = [[1, 0, 0],
   [0, 1, 0],
   [0, 0, 1]]
]

The texture data is:

  data     ∂        TensorXf BitmapTexture
  to_uv             ScalarTransform3f BitmapTexture

Could you please tell me how to properly optimize brdf texture in this case?
Thank you

[njroussel]

Hi @osylum

Sounds like you might not be updating the scene parameter after you've created the optimizer:

params.update(opt);

Basically the scene parameters need to be aware that the optimizer is tracking its gradients.

This line dr.enable_grad(params['data']) is unnecessary, the optimizer will do it for you.

[osylum]

I added both params.update(opt) and scene_params.update(opt) after setting the optimizer parameters with the same error:

TypeError: backward_from(): the argument does not depend on the input variable(s) being differentiated. Raising an exception since this is usually indicative of a bug (for example, you may have forgotten to call dr.enable_grad(..)). If this is expected behavior, skip the call to backward_from(..) if ek.grad_enabled(..) returns False.

[njroussel]

(Please format your answers with code blocks for error messages, and answer in the same thread. It makes it easier to read and to follow for myself and anyone else coming across this thread).

I don't understand why you have duplicate scene parameters. Please follow the tutorials, they serve as a great starting point.

I put together a minimal texture optimisation script here:

import mitsuba as mi
import drjit as dr

mi.set_variant('cuda_ad_rgb')

# Low resolution reference
reference_tex = mi.Bitmap('mitsuba3/resources/data/common/textures/flower.bmp')

scene_description = mi.cornell_box()
scene_description['integrator'] = mi.load_dict({
    'type': 'prb',
})
scene_description['red-wall']['bsdf'] = mi.load_dict({
    'type': 'diffuse',
    'reflectance': {
        'type': 'bitmap',
        'bitmap': reference_tex
    }
})

scene = mi.load_dict(scene_description)
ref = mi.render(scene, spp=2048)

key = 'red-wall.bsdf.reflectance.data'
params = mi.traverse(scene)
params.keep([key])

# Start with almost all black texture
start_params = dr.full(mi.TensorXf, 0.05, mi.TensorXf(reference_tex).shape)
params[key] = start_params
params.update()
starting = mi.render(scene)

opt = mi.ad.Adam(lr=2., params=params) # you can also ignore the `params`  kwarg and use opt[key] = params[key]
params.update(opt)

for i in range(10):
    img = mi.render(scene, params, spp=32)
    loss = dr.mean(dr.sqr(img - ref))

    dr.backward(loss)
    opt.step()
    opt[key] = dr.clamp(opt[key], 0., 1.0)
    params.update(opt)

optimized = mi.render(scene, spp=2048)

mi.Bitmap(ref).write('ref.exr')
mi.Bitmap(starting).write('starting.exr')
mi.Bitmap(optimized).write('optimized.exr')

[osylum]

sure, I upated the style for the errors.

I had followed the tutorials. But in my case, there is a difference as compared to the tutorials or your example.

In your example, mi.traverse(scene), the texture to update is there:

red-wall.bsdf.reflectance.data     ∂        TensorXf BitmapTexture

In my case, when using scene_params = mi.traverse(scene), the parameter I want to access shows up as:

object.bsdf.nested_bsdf.rho_d        ∂        Texture BSDF

and need to further traverse this parameter to access the data, as param = mi.traverse(scene_params['object.bsdf.nested_bsdf.rho_d']):

data     ∂        TensorXf BitmapTexture

[njroussel]

Oh I see!
I'm guessing you're using a custom BSDF then?

If so, you should be using put_object in its traverse() method rather than put_parameter. This will recursively traverse the objects which should give you directly access to the tensor rather than having to traverse again.

In short, in this process we basically want the scene to point to the same object (parameter to be differentiated) as params and the optimizer. I believe having a second traverse in your case is breaking this fragile reference logic 🙁 .

[osylum]

Great! now I can see it in the scene_params list:

object.bsdf.nested_bsdf.rho_d.data         ∂        TensorXf BitmapTexture

and do not get an error during the call to dr.backward.
yep, I am using a custom BSDF (see first comment 😉).