How to release the memory in the loop of ray sampling

[boringfish]

Hi, everyone. I'm trying to create a radiosity dataset. Specifically, I load a scene, transform the scene parameters each time through a for loop, shoot a ray and calculate the radiosity at the intersection point. Here is my code:

    mi.set_variant('cuda_ad_rgb')
    # Load a scene
    scene = mi.load_file(
        "C:/Users/China/Downloads/mitsuba-data-master/mitsuba-data-master/docs/scenes/medium_testheterogeneous_sss.xml")
    sampler = mi.load_dict({'type': 'independent', 'sample_count': 16})
    integrator = mi.load_dict({'type': 'path'})
    offset = mi.Vector3f(0.5, 0.5, 0.5)

    # todo: g,albedo,sigma_t_scale
    params = mi.traverse(scene)
    to_world = params['Area_002-light.to_world']
    dest_vol = params['medium1.sigma_t.data'].numpy()
    albedo_vol = params['medium1.albedo.data'].numpy()

    for i in trange(200000):
        scale = np.random.uniform(0.1, 1000)
        params['medium1.scale'] = scale

        g = np.random.uniform(0.0, 1)
        params['medium1.phase_function.g'] = g

        albedo_scale = np.random.uniform(0.01, 10)
        params['medium1.albedo.data'] = mi.TensorXf(np.clip(albedo_vol * albedo_scale, 0, 1))


        rotation_x = mi.Transform4f.rotate(axis=[1, 0, 0], angle=np.random.uniform(0, 360))
        rotation_y = mi.Transform4f.rotate(axis=[0, 1, 0], angle=np.random.uniform(0, 360))
        rotation_z = mi.Transform4f.rotate(axis=[0, 0, 1], angle=np.random.uniform(0, 360))
        new_to_world = rotation_z @ rotation_y @ rotation_x @ to_world
        params['Area_002-light.to_world'] = new_to_world
        params.update()

        origin = sample_spherical() + offset
        # normalize direction
        normalized_direction = dr.normalize(sample_spherical(radius=0.5) + offset - origin)
        ray = mi.Ray3f(origin, normalized_direction)
        si = scene.ray_intersect(ray)

        while (si.t == np.inf):
            origin = sample_spherical() + offset
            normalized_direction = dr.normalize(sample_spherical(radius=0.5) + offset - origin)
            ray = mi.Ray3f(origin, normalized_direction)
            si = scene.ray_intersect(ray)

        # --- ray batch，with the same origin and direction

        ray_count = 4096 
        origins = dr.tile(origin, ray_count)
        directions = dr.tile(normalized_direction, ray_count)
        rays = mi.Ray3f(origins, directions)
        random_number = np.random.randint(0, 10000)
        sampler.seed(random_number, wavefront_size=ray_count)

        radiance_batch = integrator.sample(scene, sampler, rays)

        mean_radiance = np.mean(np.array(radiance_batch[0]), axis=0)
        sampler.schedule_state()  # <- important!

        print(mean_radiance)

However, I found that the memory occupied by the Python program would increase rapidly in each loop and eventually reach the upper limit of the system. A similar issue can be seen in discussion 1171. I tried to use cleanup() in each loop but the program went wrong. I guess the reason is that the scene is shared between each loop. Therefore, is there any good way to release the memory used in each loop? Thanks a lot for your help.

[njroussel]

Hi @boringfish

Could you provide us with a full script ? We're currently missing the sample_spherical function in order to run this ourselves.

I didn't spot anything obvious. However, one thing that came to mind is the following issue: mitsuba-renderer/drjit#137
Are you running out of VRAM or RAM ?

[boringfish]

Sure. The sample_spherical is a simple function that samples coordinate points on the sphere. Here is my full code:

import numpy as np
import mitsuba as mi
from tqdm import trange
import drjit as dr

def sample_spherical(radius=1.5):
    # Step 1: Generate two uniform random numbers
    u = np.random.rand()
    v = np.random.rand()

    # Step 2: Convert to spherical coordinates
    theta = 2 * np.pi * u  # Azimuthal angle
    phi = np.arccos(2 * v - 1)  # Polar angle

    # Step 3: Convert to Cartesian coordinates
    x = radius * np.sin(phi) * np.cos(theta)
    y = radius * np.sin(phi) * np.sin(theta)
    z = radius * np.cos(phi)

    return mi.Point3f(x, y, z)

# Press the green button in the gutter to run the script.
if __name__ == '__main__':
    mitsuba_data_path = "C:/Users/China/Downloads/mitsuba-data-master/"

    mi.set_variant('cuda_ad_rgb')
    # mi.set_variant('scalar_rgb')
    # Load a scene
    scene = mi.load_file(
        "C:/Users/China/Downloads/mitsuba-data-master/mitsuba-data-master/docs/scenes/medium_testheterogeneous_sss.xml")
    sampler = mi.load_dict({'type': 'independent', 'sample_count': 16})
    integrator = mi.load_dict({'type': 'path'})

    offset = mi.Vector3f(0.5, 0.5, 0.5)

    # todo: g,albedo,sigma_t_scale
    params = mi.traverse(scene)
    to_world = params['Area_002-light.to_world']
    dest_vol = params['medium1.sigma_t.data'].numpy()
    albedo_vol = params['medium1.albedo.data'].numpy()

    for i in trange(200000):
        scale = np.random.uniform(0.1, 1000)
        params['medium1.scale'] = scale

        g = np.random.uniform(0.0, 1)
        params['medium1.phase_function.g'] = g

        albedo_scale = np.random.uniform(0.01, 10)
        params['medium1.albedo.data'] = mi.TensorXf(np.clip(albedo_vol * albedo_scale, 0, 1))


        rotation_x = mi.Transform4f.rotate(axis=[1, 0, 0], angle=np.random.uniform(0, 360))
        rotation_y = mi.Transform4f.rotate(axis=[0, 1, 0], angle=np.random.uniform(0, 360))
        rotation_z = mi.Transform4f.rotate(axis=[0, 0, 1], angle=np.random.uniform(0, 360))
        new_to_world = rotation_z @ rotation_y @ rotation_x @ to_world
        params['Area_002-light.to_world'] = new_to_world
        params.update()

        origin = sample_spherical() + offset
        # normalize direction
        normalized_direction = dr.normalize(sample_spherical(radius=0.5) + offset - origin)
        ray = mi.Ray3f(origin, normalized_direction)
        si = scene.ray_intersect(ray)

        while (si.t == np.inf):
            origin = sample_spherical() + offset
            normalized_direction = dr.normalize(sample_spherical(radius=0.5) + offset - origin)
            ray = mi.Ray3f(origin, normalized_direction)
            si = scene.ray_intersect(ray)

        # --- ray batch，with the same origin and direction

        ray_count = 4096 * 256
        origins = dr.tile(origin, ray_count)
        directions = dr.tile(normalized_direction, ray_count)
        rays = mi.Ray3f(origins, directions)
        random_number = np.random.randint(0, 10000)
        sampler.seed(random_number, wavefront_size=ray_count)

        radiance_batch = integrator.sample(scene, sampler, rays)

        mean_radiance = np.mean(np.array(radiance_batch[0]), axis=0)
        sampler.schedule_state()  # <- important!

        print(mean_radiance)

BY the way, I run out of RAM. In the Windows Task Manager, I can see that the memory is growing rapidly during the loop.