Unexpected Error when Optimizing over Directional Light

[jensenzhoujh]

Hi @njroussel,

I am trying to use mitsuba to optimize the intensity of two light sources, with one being constant enviroment light and the other being directional light. The code snippet is as below and can run sucessfully for both keys either const_env_light.radiance.value or directional_light.irradiance.value.

import drjit as dr
import mitsuba as mi
mi.set_variant('cuda_ad_rgb')

import torch
import torch.nn as nn
import torch.nn.functional as F

from matplotlib import pyplot as plt
from mitsuba.scalar_rgb import Transform4f as T

filenames = [
    'mitsuba/scenes/textures/flower_photo_downscale.jpeg',
]

textures = [
    mi.TensorXf(mi.Bitmap(f).convert(mi.Bitmap.PixelFormat.RGB, mi.Struct.Type.Float32))
    for f in filenames
]

res = dr.shape(textures[0])[0]

scene = mi.load_dict({
    'type': 'scene',
    'integrator': {
        'type': 'direct_reparam',
    },
    'sensor':  {
        'type': 'perspective',
        'to_world': T.look_at(
                        origin=(0, 0, -2),
                        target=(0, 0, 0),
                        up=(0, -1, 0)
                    ),
        'fov': 60,
        'film': {
            'type': 'hdrfilm',
            'width':  res,
            'height': res,
            'sample_border': True,
        },
    },
    'textured_plane': {
        'type': 'rectangle',
        'to_world': T.scale(1.2),
        'bsdf': {
            'type': 'twosided',
            'nested': {
                'type': 'diffuse',
                'reflectance': {
                    'type': 'bitmap',
                    'filename': filenames[0]
                },
            }
        }
    },
    'const_env_light': {
        'type': 'constant',
        'radiance': {
            'type': 'rgb',
            'value': 0.5,
        }
    },
    'directional_light': {
        'type': 'directional',
        'direction': [0, 0, 1],
        'irradiance': {
            'type': 'rgb',
            'value': 5,
        }
    },
})

params = mi.traverse(scene)
print(params)

optim_params = {
    'const_env_light.radiance.value': torch.full((1, 3), 0.5, device='cuda', requires_grad=True),
    'directional_light.irradiance.value': torch.full((1, 3), 5.0, device='cuda', requires_grad=True),
}

# key = 'const_env_light.radiance.value'
key = 'directional_light.irradiance.value'

dr.enable_grad(params[key])

@dr.wrap_ad(source='torch', target='drjit')
def render(value, spp=256):
    params[key] = dr.unravel(mi.Color3f, value)
    params.update()
    return mi.render(scene, params, spp=spp)

optimizer = torch.optim.Adam(optim_params.values(), lr=0.0002)
loss_fn = nn.L1Loss()

# Optimization hyper-parameters
iteration_count = 100
spp = 4

train_losses = []
for i in range(iteration_count):
    loss_accum = 0
    optimizer.zero_grad()
    rendered_img = render(optim_params[key], spp=spp)
    loss = loss_fn(rendered_img, textures[0].torch())
    loss.backward()
    loss_accum += loss.item()
    optimizer.step()
    train_losses.append(loss_accum)
    print(f'Training iteration {i+1}/{iteration_count}, loss: {train_losses[-1]}', end='\r')

However, when I change the integrator from direct_reparam to aov with a nested direct_reparam.

    'integrator': {
        'type': 'aov',
        'aovs': 'shape_idx:shape_index',
        'image': {
            'type': 'direct_reparam',
        }
    },

and accordingly change the predication from rendered_img to rendered_img[..., 4:7] for loss computation, the following bug occurs when I am trying to optimize over directional_light.irradiance.value:

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 dr.grad_enabled(..) returns False.

Strangely, everything is ok for optimization over const_env_light.radiance.value.

I've also experimented with aov with direct nested:

    'integrator': {
        'type': 'aov',
        'aovs': 'shape_idx:shape_index',
        'image': {
            'type': 'direct',
        }
    },

The optimization is fine for both const_env_light.radiance.value and directional_light.irradiance.value. I also tried using pure mitsuba optimization instead of insert it into PyTorch optimization and the phenonmenon is the same. So this should not be a problem occuring expectedly from communication between the two.

Since in my case I need aov to compute shape_idx and direct_reparam to account for visibility discontinuties (as I am also optimizing the vertice positions of the object), do you have any idea of how I can cope with this?

[njroussel]

Hi @shallowtoil

Could you send the snippet that fails without PyTorch?

[jensenzhoujh]

Hi @njroussel,

import drjit as dr
import mitsuba as mi
mi.set_variant('cuda_ad_rgb')
from mitsuba.scalar_rgb import Transform4f as T

filenames = [
    'mitsuba/scenes/textures/flower_photo_downscale.jpeg',
]

textures = [
    mi.TensorXf(mi.Bitmap(f).convert(mi.Bitmap.PixelFormat.RGB, mi.Struct.Type.Float32))
    for f in filenames
]

res = dr.shape(textures[0])[0]

scene = mi.load_dict({
    'type': 'scene',
    'integrator': {
        'type': 'aov',
        'aovs': 'shape_idx:shape_index',
        'image': {
            'type': 'direct_reparam',
        }
    },
    'sensor':  {
        'type': 'perspective',
        'to_world': T.look_at(
                        origin=(0, 0, -2),
                        target=(0, 0, 0),
                        up=(0, -1, 0)
                    ),
        'fov': 60,
        'film': {
            'type': 'hdrfilm',
            'width':  res,
            'height': res,
            'sample_border': True,
        },
    },
    'textured_plane': {
        'type': 'rectangle',
        'to_world': T.scale(1.2),
        'bsdf': {
            'type': 'twosided',
            'nested': {
                'type': 'diffuse',
                'reflectance': {
                    'type': 'bitmap',
                    'filename': filenames[0]
                },
            }
        }
    },
    'const_env_light': {
        'type': 'constant',
        'radiance': {
            'type': 'rgb',
            'value': 0.5,
        }
    },
    'directional_light': {
        'type': 'directional',
        'direction': [0, 0, 1],
        'irradiance': {
            'type': 'rgb',
            'value': 5,
        }
    },
})

params = mi.traverse(scene)
print(params)

# key = 'const_env_light.radiance.value'
key = 'directional_light.irradiance.value'

dr.enable_grad(params[key])

optimizer = mi.ad.Adam(lr=0.025)
optimizer[key] = params[key]
params.update(optimizer)
loss_fn = lambda x, y: dr.mean(dr.sqr(x - y))

# Optimization hyper-parameters
iteration_count = 100
spp = 4

train_losses = []
for i in range(iteration_count):
    loss_accum = 0
    rendered_img = mi.render(scene, params, spp=spp)
    loss = loss_fn(rendered_img[..., :3], textures[0])
    dr.backward(loss)
    loss_accum += loss
    optimizer.step()
    train_losses.append(loss_accum)
    params.update(optimizer)
    print(f'Training iteration {i+1}/{iteration_count}, loss: {train_losses[-1]}', end='\r')

[njroussel]

Hi @shallowtoil

Great, thank you! I can reproduce this.

From a quick look at it, it looks like a bug and I believe that I understand where it's coming from. Basically, the aov integrator is not calling the appropriate methods of its nested integrator for it to be differentiated. It works fine if you nest a direct integrator because here its only leveraging the AD system, whereas with direct_reparam their is an adjoint definition of the backward pass.

The only workaround I can offer to you is to use two different render statements: one that doesn't use the aov integrator and one that does.

[jensenzhoujh]

Hi @njroussel,

Will you consider fixing this bug? I can use the workaround temporarily. Since rendering twice will double the cost, it would be great if the bug can be fixed anytime soon.

[jensenzhoujh]

I've tried using two render() calls but the training will become infeasibly slow.

[njroussel]

There's no plan to work on this anytime soon. It would either require a lot of refactoring of the exisiting AOV integrator, or rewriting it entirely in Python.

The cost for the extra rendering shouldn't be that large. Basically the one that uses the AOV integrator doesn't need to do a full-depth rendering. All you want is the shape_index so this should be enough

integrator_aov = mi.load_dict({ 'type' : 'aov', 'aovs': 'si:shape_index'})
mi.render(scene, integrator=integrator_aov, spp=1) 

This is basically width x height more ray intersections which should be marginal.

[jensenzhoujh]

Indeed mi.render(scene, integrator=integrator_aov, spp=1) doesn't add too much of extra const. But using

        'integrator': {
            'type': 'aov',
            'aovs': 'shape_idx:shape_index',
            'image': {
                'type': 'direct_reparam',
            }
        },

is much faster (1s/iter) than

        'integrator': {
            'type': 'direct_reparam',
        },

which takes approximately (28s/iter). As comparison, using

        'integrator': {
            'type': 'direct',
        },

takes about 2s/iter. Do you have any idea what happened here? Why calling direct_reparam alone is strangely slow? I use a spp=128 otherwise there will be too much of noise in the rendered image.

[njroussel]

The direct and direct_reparam integrators do not do the same thing in their backward pass. As I mentioned before direct_reparam has its own "definition/implementation" for it such that it can handle discontinuities. This explains why direct_reparam is considerably slower than direct. You should read more about this in the first section of this tutorial

Now why is it faster again when wrapped with the aov? Again, the aov is not calling the appropriate methods during the backward pass hence it is not making use of the direct_reparam's backward mode but producing something similar to direct.

You don't need to nest an integrator in the aov. In my previous message it is an integrator that only returns the shape_index. Your backward() call should only be made with the direct/direct_reparam outputs -- the shape_index isn't differentiable anyway.

[njroussel]

I should add that you might indeed want a higher spp for your direct_reparam render, but the AOV shape_index can be obtained with just 1 spp.