Question on forward mode for geometry parameters

[joyDeng]

Summary

Question on differentiating the geometries with forward mode in python with mitsuba.set_variant('gpu_autodiff_rgb') .

System configuration

• Platform: Ubuntu 19.10
• Compiler: clang
• Python version: 3.7.4
• Mitsuba 2 version: 2.2.1
• Compiled variants: "scalar_rgb", "scalar_spectral", "packet_rgb", "packet_spectral", "gpu_autodiff_spectral", "gpu_autodiff_rgb"

Description

Basically the quantity I want is d I / d x , where I is pixel values and x is the vertex positions. I took the code from test_mesh.py and the example code from documentation where it computes gradient of reflectance on the wall with forward mode. However, when I use ek.set_gradient(param, perturb, backward=false) and then call ek.forward(param), it always gives me the same value whatever perturb I set. For example, I have param which is the vertex position buffer = [x, y, z] and I want to compute d I / dx, dI / dy, dI / dz respectively. I did the followings to compute the d I / dx:

1. I set img = render(sence,spp=4), then I assign perturb = [1, 0, 0] ,
2. call ek.set_gradient(param, perturb, backward=false) ,
3. call ek.forward(param),
4. and dI / dx = ek.gradient(img).

Then I compute dI / dy by reset the perturb to [0, 1, 0] and run the setp 2 to 4 again. But when I visualize the results, it turns out that d I / dy , dI / dx and d I / dz have the same values. I strongly suspect that I misunderstood something, but this is what I could infer from the material I had. If anyone could provide any insights, or any example code of differentiating vertices with forward mode, it will be really appreciated! updates_____ However, I could compute the individual partial derivative gradient now by re-render the image every time before I set the gradient. Is it the correct way to do that? it sounds not efficient because when I try to render it multiple times with spp=4 it gives me an memory error information:

cuda_check(): runtime API error = 0002 "cudaErrorMemoryAllocation" in ../ext/enoki/src/cuda/horiz.cu:59. terminating with uncaught exception of type std::runtime_error: CUDABackend: referenced unknown variable 457955 Aborted

Reproduce the bug

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

import numpy as np
from mitsuba.core import Color3f, Float, Thread, Ray3f, Vector3f, Vector2f, UInt32, Transform4f
from mitsuba.core.xml import load_file
from mitsuba.python.util import traverse
from mitsuba.python.autodiff import render, write_bitmap
from enoki.cuda_autodiff import Float32 as FloatD


filepath = "./scenes/cbox/cbox.xml“
key = “cbox_smallbox.vertex_positions_bu”'

Thread.thread().file_resolver().append("cbox")
scene = load_file(filepath)
crop_size = scene.sensors()[0].film().crop_size()
    
params = traverse(scene)
vertex_ref = np.array(params[key]).tolist()
params.keep([key])
ek.set_requires_gradient(params[key])
params.set_dirty(key)
params.update()

img = render(scene, 4)
partial_grads = np.zeros((0, len(img)))

# compute a partial derivative for each element in position buf
for i in range(len(vertex_ref)):

    # set the perturb of the paramter w.r.t whose partial derivative we want to 1.0
    perturb = (np.zeros(len(vertex_ref))).tolist()
    perturb[i] = 1

    # if I delete the line below, every loop gives me the same value
    # but spp=8 gives me memory error, setting spp=4 will make it run.
    # Though, I still not sure whether it is the noise or the partial derivatives differences
    img = render(scene, 8) 
    
    ek.set_gradient(params[key], perturb, backward=False)
    ek.forward(params[key])
    grad = ek.gradient(img)

    # this might contain negative value, but can still visualize in tev
    write_bitmap('vertex_grad_{:03}.exr'.format(i), grad, crop_size)

    gradnp = np.array(grad)
    partial_grads = np.concatenate((partial_grads, [gradnp]), axis=0)

Scene file:

Cornell box scene

[Speierers]

Hi @joyDeng,

Forward-mode autodiff can only compute image gradient w.r.t. a single scene parameter at a time. This is extremely inefficient in the case of meshes as you will need to perform this propagation 3 x vertex_count times. But the computation (e.g. render() should only be needed once)

That being said, it should work in practice. Looking at your code above, I think those two lines are the problem:

ek.set_gradient(params[key], perturb, backward=False)
ek.forward(params[key])

You only need to call ek.set_gradient() once before your computations, so that enoki build the AD graph for those variables. It is already done (right before calling params.set_dirty(key) which is good. And the first call to render() occurs after that which is correct as well.

ek.forward(params[key]) isn't correct as params[key] is an array and not a single "scalar" variable. I suppose in this case enoki always forward the first scalar variable in this array that has gradients enabled.

In theory in the for loop, all you need is to call ek.forward for a single "scalar" variable and then query the respective image gradients. It should look something like this:

ek.forward(params[key][i])
grad = ek.gradient(img)

but I doubt this will work as params[key][i] isn't even an enoki array.

One solution for this would be to create a list of individual Float(0.0) variables (3 per vertices), and add those to the mesh.vertex_buffer (e.g. using an addition). Those should all have set_enable_gradient(true) and will become the entry points of the AD graph. After the computation done, you can simply call ek.forward(var_list[i]) to propagate the gradients for each one of those variables.

I hope this helps.

[joyDeng]

Hi @Speierers

Thank you for the reply. I am trying this method, however I am not very clear how to add the float to vertex buffer, what I am doing now is:

params[key][i] = var_list[i]

where var_list is a list of FloatD each related to one value in vertex buffer

but it gives me segmentation fault when I update the params, it seems that the parameter type is not correct,

then I tried: params[key][i] = params[key][i] * 0 + var_list[i]

But this operation overwrites all the value in the params[key] with var_list[i] and still give segmentation fault.

[Speierers]

Actually my workaround will indeed not work as the vertex position buffer is a single enoki array.

Could you try something like this (untested):

for i in range(len(params[key])):
    vertex_pos_grad = Float.zero(len(params[key]))
    ek.scatter(vertex_pos_grad, i, 1)
    ek.set_gradient(params[key], vertex_pos_grad)
    Float.forward()
    grad = ek.gradient(img)
    ...

This is very similar to what you have at the beginning, except that here we are directly calling Float.forward() instead of ek.foward(my_var)` as we manually set the gradients of the leaf node.

[joyDeng]

Thank you!

I got this:

ek.set_gradient(sc.params[sc.key], vertex_pos_grad, backward=False) TypeError: set_gradient(): incompatible function arguments

by the way, should I add backward=false in this function?

[Speierers]

Yes you will need backward=False.

Maybe could you try to ek.detach() the vertex_pos_grad variable when passing it to ek.set_gradient()?

[joyDeng]

Thank you Speierers! The code runs but the gradient it gives me are all the same no matter what I put into set_gradient. Is it always forward the first scalar variable in this array that has gradients enabled? Or is it possible to compute partial derivative matrix with backward mode?

[Speierers]

It will traverse the AD graph in the forward direction and propagate any available gradients. So as you only set non-zero gradient to a single vertex component, this you give you the right result. Could you maybe try with a simpler scene (e.g. just a single triangle) and print the gradients for the vertex positions after calling ek.set_gradient() to ensure this works properly?

Backward mode will only propagate gradients (in backward direction) starting at a single node (e.g. a loss value on the entire image, a single pixel value, ...). You won't be able to reconstruct "gradient images" this way unfortunately.