How to using the modified intermediate gradients during optimization?

[ryukang]

Hello,

I'd like to modify intermediate gradient, and then recalculate parameter's gradient value.

For example, in in gradient-based optimization tutorial, I was able to get the gradient value (dL/dI) of the image using backward(loss, dr.ADFlag.ClearEdges).

After modifying the gradient value obtained here, I want to recalculate the parameter's gradient value using the modified dL/dI. (e.g., modified dL/dI * dI/d(opt[key]))

loss = mse(image)
dr.backward(loss, dr.ADFlag.ClearEdges)

modified_value = dr.grad(image) * (...)
dr.set_grad(image, modified_value)

???

I knew that the gradient for the image could be modified using dr.set_grad(image, ...), but I couldn't find a way to recalculate using the modified gradient.

Can anyone point me in the right direction here?

Thanks!

[njroussel]

Hi @ryukang

There are fine controls available in Dr.Jit for gradient propagation. In particular, you might want to have a look at dr.backward_to.
You could do:

loss = mse(image)
dr.backward_to(image)
dr.set_grad(image, modified_value)
dr.backward(image)

The backward_to call will stop the propagation at the image variable.

EDIT:
Sorry, I read through that too quickly. In the case where your modified_value depends on the gradients of your parameters for a "typical" backpropagation, you'd want to do this:

loss = mse(image)
dr.backward_to(image)  # Differentiate through loss function
dr.backward(image, dr.ADFlag.ClearInterior) # Propagate to parameters without destroying the graph edges

 # Prepare new modified gradient
modified_value = some_function(dr.grad(image), dr.grad(opt[key]))
dr.set_grad(image, modified_value)


dr.backward(image) 

[ryukang]

Hi, @njroussel ,
Thank you so much for the quick reply!
When I proceeded as suggested, it didn't go well.

First,

loss = mse(image)
dr.backward_to(image)  
dr.backward(image, dr.ADFlag.ClearInterior) 

In this part, dr.grad(opt[key]) is well obtained, but when I print dr.grad(image), TensorXf(1.0) is output.

But what I want is the original dL/dI value, not the value of 1.

Second,

modified_value = some_function(dr.grad(image), dr.grad(opt[key]))
dr.set_grad(image, modified_value)

dr.set_grad(opt[key], 0) # Clear current gradient
dr.backward(image)

In this part, I simply multiplied the dr.grad(image) by 0.001 and then performed backward.

dr.set_grad(image, dr.grad(image) * 0.001)
dr.set_grad(opt[key], 0) # Clear current gradient
dr.backward(image)

However, after doing backward(image), the value of dr.grad(opt[key]) almost no change.

When I studied more, it seems that when calling backward(arg), the gradient value in 'arg' is initialized to 1.

.
.
.

I think I asked a wrong question. so I ask you a question again.
I would like to obtain the original dL/dI value, modify this value (e.g., dL/dI * 0.1), and recalculate the gradient of loss with respect to parameters (dr.grad(opt[key])) using modified gradient. How would I go about doing this?

Thanks for your help!

[njroussel]

Sorry, I responded too quickly.

You are correct. dr.backward and dr.forward are just aliases for, respectively, dr.backward_from and dr.forward_from. Both functions set the variable's gradient to 1 before propagating.

Here's my corrected suggestion:

import mitsuba as mi
import drjit as dr

mi.set_variant('cuda_ad_rgb')

scene_description = mi.cornell_box()
scene_description['integrator']['type'] = 'prb'
scene = mi.load_dict(scene_description)

params = mi.traverse(scene)
key = 'red.reflectance.value'

prev_value = params[key]
new_value = mi.Color3f([0, 0, 1])

params[key] = new_value
params.update()
ref_img = mi.render(scene, params=params, seed=0xDEADBEEF)

params[key] = prev_value
dr.enable_grad(params[key])
params.update()
img = mi.render(scene, params=params, seed=0xDEADBEEF)

loss = dr.sum(dr.sqr(img - ref_img))

dr.set_grad(loss, 1) # Initialize gradient
dr.backward_to(img) # Stop propagation at the image
dr.set_grad(img, dr.grad(img) * 100) # Modify gradient of the image
dr.backward_to(params[key]) # Propagate from the image to parameters

[ryukang]

Thanks! I will try it.