Unexpected error using mi.load_dict in loop

[jensenzhoujh]

Hi @njroussel ,

I am trying to instantiate the mesh and scene using mi.load_dict inside a loop. Basically, I have a data loader sampling a bunch of meshes and for each mesh I will optimize the scene parameters. When using integrated optimization with PyTorch the following error occurs:

Info: Pipeline statistics
        module(s)                            :     2
        entry function(s)                    :    25
        trace call(s)                        :     3
        continuation callable call(s)        :     0
        direct callable call(s)              :     9
        basic block(s) in entry functions    :   503
        instruction(s) in entry functions    : 13414
        non-entry function(s)                :     0
        basic block(s) in non-entry functions:     0
        instruction(s) in non-entry functions:     0
        debug information                    :    no



Critical Dr.Jit compiler failure: jit_optix_check(): API error 7251 (OPTIX_ERROR_PIPELINE_LINK_ERROR): "Pipeline link error" in /project/ext/drjit-core/src/optix_core.cpp:392%

The following code should be able to re-produce the problem.

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

for _ in range(2):
    filenames = [
        '../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': 'prb_reparam',
                'max_depth': 2,
                'reparam_max_depth': 2,
            }
        },
        # '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 = 128

    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[..., :3], 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')

When I'm not using the PyTorch integrated version, it works perfectly fine (see below).

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

for _ in range(2):
    filenames = [
        '../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',
        # },
        '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 = 128

    start_time = time.time()
    train_losses = []
    for i in range(iteration_count):
        loss_accum = 0
        rendered_img = mi.render(scene, params, spp=spp)
        # print(optim_params[key])
        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')
    end_time = time.time()
    print(f'Training time: {end_time - start_time:.2f}s')
    del scene, params

Do you have any workaround to resolve this temporarily? I don't want to loop the dataloader outside the file since i need to load a pre-trained model to supervise the update of scene parameters which takes a fair amount of time especially when I evaluate on a large dataset of 3D meshes.

[njroussel]

Hi @shallowtoil

This seems like a bug.
I do need some clarifications though. So, if you have the load_dict outside of the loop it seems to work?

I tried running this myself and am able reproduce the error. Actually, it's not exactly the same error but still an OptiX error. What surprises me is that the loss can be computed without any issues, it's the backward() that is triggering the bug here. Meaning that the OptiX scene becomes invalidated not between the loop iterations but during the the second iteration.

[jensenzhoujh]

Yes. There will be no bug if load_dict is outside the loop.

[qiangdadi]

Hi,
Has the bug been fixed? I encounter an error similar to this one.

[merlinND]

Hello @qiangdadi,

This is a post from 2023 and much has changed since then. Please open a new Issue with all details about your setup (OS, GPU, driver version, etc), as well as full code and data used to reproduce your problem.

[Speierers]

A few questions:

1. Does this bug also occur with the llvm variants?
2. Does it help to add dr.sync_thread() at the end of the outer loop? Basically to make sure that all kernel evalutations are finished before moving on to the next iteration. (also maybe try to print the optimized parameter at the end of the outer loop body)