[🐛 bug report] Gradient value is empty with free_graph=False

[pjessesco]

Summary

Hi, I'm trying to do a differentiable rendering without free_graph option. It looks like if the option is enabled, the visited node is erased after used to calculate gradient while traversing the graph.

My goal is to specify nodes' gradient in their label in graph (using ek.graphviz), so I don't want to let my graph be reduced while backpropagation.

However I couldn't optimize parameter with free_graph=False... thanks for your reply.

System configuration

master b92ddc2
Ubuntu
CUDA 11.2

Steps to reproduce

In invert_bunny.py example, fix as below

ek.backward(ob_val, free_graph=False)
...
...
print('Iteration %03i: error=%g' % (it, err_ref[0]), ek.gradient(params['my_envmap.data']))

[Speierers]

Hi @pjessesco ,

What's do you mean by "can't optimize"? Are the resulting gradients wrong?

[pjessesco]

Hi @Speierers , sorry for ambiguous wording.

free_graph = True

free_graph = False

I printed as

print('Iteration %03i: error=%g' % (it, err_ref[0]), ek.gradient(params['my_envmap.data']))

, gradient value is empty if the option is false.

[Speierers]

That is likely a bug in the enoki AD backend. Unfortunately we are working hard on the upcoming release of the library I won't probably have the time to look into this in the near future.

In any case, could you send me the whole python snippet to reproduce this, I could take a brief look to see if there is anything obvious.

[pjessesco]

Here is the script, it's almost the same with a bunny example except for the option.

invert_bunny.py

# Advanced inverse rendering example: render a bunny reference image,
# then replace one of the scene parameters and try to recover it using
# differentiable rendering and gradient-based optimization.

import enoki as ek
import mitsuba
mitsuba.set_variant('gpu_autodiff_rgb')

from mitsuba.core import Float, Thread
from mitsuba.core.xml import load_file
from mitsuba.python.util import traverse
from mitsuba.python.autodiff import render, write_bitmap, Adam
import time

# Load example scene
Thread.thread().file_resolver().append('bunny')
scene = load_file('bunny/bunny.xml')

# Find differentiable scene parameters
params = traverse(scene)

# Make a backup copy
param_res = params['my_envmap.resolution']
param_ref = Float(params['my_envmap.data'])

# Discard all parameters except for one we want to differentiate
params.keep(['my_envmap.data'])

# Render a reference image (no derivatives used yet)
image_ref = render(scene, spp=16)
crop_size = scene.sensors()[0].film().crop_size()
write_bitmap('out_ref.png', image_ref, crop_size)

# Change to a uniform white lighting environment
params['my_envmap.data'] = ek.full(Float, 1.0, len(param_ref))
params.update()

# Construct an Adam optimizer that will adjust the parameters 'params'
opt = Adam(params, lr=.02)

time_a = time.time()

iterations = 100
for it in range(iterations):
    # Perform a differentiable rendering of the scene
    image = render(scene, optimizer=opt, unbiased=True, spp=1)
    write_bitmap('out_%03i.png' % it, image, crop_size)
    write_bitmap('envmap_%03i.png' % it, params['my_envmap.data'],
                (param_res[0], param_res[1]))

    # Objective: MSE between 'image' and 'image_ref'
    ob_val = ek.hsum(ek.sqr(image - image_ref)) / len(image)

    # Back-propagate errors to input parameters
    # ek.backward(ob_val)
    ek.backward(ob_val, free_graph=False)

    # Optimizer: take a gradient step
    opt.step()

    # Compare iterate against ground-truth value
    err_ref = ek.hsum(ek.sqr(param_ref - params['my_envmap.data']))
    print('Iteration %03i: error=%g' % (it, err_ref[0]), ek.gradient(params['my_envmap.data']))

time_b = time.time()

print()
print('%f ms per iteration' % (((time_b - time_a) * 1000) / iterations))

[Speierers]

This looks reasonable. It is likely a bug in the enoki AD backend then.