interoperability and data exchange with Mitsuba and PyTorch

[lesphere]

Hi,

I'd like to restore some intermediate values of render_backward() to train a neural network, but I don't know how to do the data exchange with jit on.

In details, in src\python\python\ad\integrators\prb.py PRBIntegrator.sample(), there's a loop do something like ray tracing. In adjoint mode, this loop computes derivatives in a way like ray tracing. I use aovs to record some extra data. But I don't know how to restore the data of different threads together to generate a training dataset (in the context of machine learning).

Here is an example modified from the original code of Mitsuba 3:

def sample(..., aovs, ...):
    ...
    while loop(active):
        ...
        # there are positions of intersections, incoming and outgoing directions 
        # and some other values in aovs
        aovs.append(...)
        ...
    ...
def render_backward(...):
    ...
    aovs = []
    sample(..., aovs, ...)
    # How to put the data in aovs into the training dataset 
    # that can be used to train a NN later

In main.py:

mi.set_variant('cuda_ad_rgb')
for it in range(...):
    image = mi.render()
    loss = mse(image, ...)
    dr.backward(loss)
    # use the data generated in last frame to train one epoch

Suppose that type(aovs[0]) == mi.Point3f, can I achieve that by something like below?

pos_flat = dr.ravel(aovs[0])
pos_tensor = mi.TensorXf(pos_flat, ...)
dataset.insert_pos(pos_tensor)

If so, how can I get the dataset reference in render_backward()?
One possible way seems to be modifying all the related parameter lists in Python and C++ to pass the reference layer by layer. But is there any better practice?

Besides, which part of the code will be compiled into a kernel, which statements will be ignored, and how do each statement work finally?

[njroussel]

Hi @lesphere

I'm not quite sure I understand what you are trying to achieve.

There is no "data" in AOVs in render_backward. The AOVs are filled during the primal rendering in the normal render call. You can access them as part of the Tensor you get back from img_tensor = mi.render(). The first channels of the tensor are the typical rendered image and the later channels are the AOVs.

[lesphere]

Thank you for your patience.

In fact, I want to record values like positions of intersections, incoming and outgoing directions for each bounce to generate a training dataset. The value for each depth of each ray forms a sample in the dataset.

And still, I wonder which part of the code will be compiled into a kernel, which statements will be ignored. And which part will be executed in single thread, which part in multiple threads?

[njroussel]

In fact, I want to record values like positions of intersections, incoming and outgoing directions for each bounce to generate a training dataset. The value for each depth of each ray forms a sample in the dataset.

This should be possible by filling the aov field in the Integrator's sample method. However, unless you're rendering one ray/sample per pixel, the output AOVs will contain the average of all rays of that pixel. This most likely meaningless for something like the incoming/outgoing directions.
If you actually want to store this information per ray, my recommendation would be to move away from the traditional Mitsuba rendering pipeline and instead script your own utility pipeline (see this tutorial). With the traditional pipeline, it is impossible to extract per-ray information.

And still, I wonder which part of the code will be compiled into a kernel, which statements will be ignored. And which part will be executed in single thread, which part in multiple threads?
One mi.render() call will typically translate into two kernels. Only the statements that are relevant to compute the output are included in the kernel: if you've written some no-ops or some code that is computing some variable that is not an output of the current kernel it will be ignored.
A kernel is entirely vectorized. So, if you're rendering some image of resolution 256 x 256 with 32 samples per pixel the kernel will be 256 * 256 * 32 threads wide and every thread will run approximately the same code path. This is only true for JIT-enabled variants, in scalar mode there are no kernels.

[lesphere]

Thank you for you reply!

For the second question, what about dr.backward()? AFAIK, this function has no explicit output. Does this mean that the output are the gradients of the scene parameters with gradiant tracking enabled and the statements that have no effects on these gradients will be ignored?

[njroussel]

Does this mean that the output are the gradients of the scene parameters with gradiant tracking enabled and the statements that have no effects on these gradients will be ignored?

Yes, you're spot on.
You can read a bit about the internals of the AD system in drjit.traverse() and drjit.enqueue() or directly in the source code. In short, the recorded computation graph is traversed to find the sub-graph which only includes the relevant inputs/outputs.