memory leak when create new sensor

[lcc815]

Summary

memory leak when create new sensor

System configuration

System information:

CPU: AMD EPYC 7542 32-Core Processor GPU: NVIDIA GeForce RTX 4090 Python: 3.8.13 (default, Apr 19 2022, 00:53:22) [GCC 7.5.0] NVidia driver: 535.129.03 LLVM: 10.0.0

Dr.Jit: 0.4.6 Mitsuba: 3.5.2 Is custom build? False Compiled with: GNU 10.2.1 Variants: scalar_rgb scalar_spectral cuda_ad_rgb llvm_ad_rgb

Description

I build a custom sensor plugin, but the memory keeps being more when I keep running both load_dict and render (only run load_dict or render function cannot reproduce it).

Steps to reproduce

just run following scripts and see nvidia-smi:

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

class FakeCamera(mi.Sensor):
    def __init__(self, props=mi.Properties()):
        super().__init__(props)

    def sample_ray(self, time, wavelength_sample, position_sample, aperture_sample, active=True):
        wavelengths, wav_weight = self.sample_wavelengths(
            dr.zeros(mi.SurfaceInteraction3f), wavelength_sample, active
        )

        film_size = mi.Vector2f(self.film().size())
        image_sample = position_sample * film_size

        x_cam = image_sample.x  # whatever the code is...
        y_cam = image_sample.y
        z_cam = image_sample.x

        camera_sample = mi.Point3f(
            x_cam,
            y_cam,
            z_cam
        )
        # Create ray direction
        d = self.world_transform() @ camera_sample
        o = self.world_transform().translation()  # Ray origin
        d = dr.normalize(d - o)

        return mi.Ray3f(o, d, time, wavelengths), wav_weight

mi.register_sensor("fakecamera", lambda props: FakeCamera(props))

if __name__ == "__main__":
    import time

    for i in range(1000):
        scene = mi.load_file("scenes/cbox.xml")
        sensor = mi.load_dict({
            'type': 'fakecamera',
            'to_world': mi.ScalarTransform4f.look_at(
                        origin=[0, 0., 1],
                        target=[0, 0, 0],
                        up=[0, -1, 0]
                    ),
            'sampler': {
                'type': 'independent',
                'sample_count': 16
            },
            'film': {
                'type': 'hdrfilm',
                'width': 1920,
                'height': 1080,
                'pixel_format': 'rgb',
                'component_format': 'float32'
            }
        })

        t1 = time.time()
        img = mi.render(scene, sensor=sensor, spp=16)  
        # print(f"Rendering costs {time.time() - t1:.3f}s")

[njroussel]

Hi @lcc815

This is a known issue: any custom Python plugin will never be properly freed. For various reasons, this is an internal limitation. In this example, although the FakeCamera doesn't have any local storage, it is still holding a reference to the Film which will therefore also not be cleaned up and it holds at least a buffer of the size of your rendered image.

The master branch has a different implementation for this feature, which will no longer leak memory.

As a workaround, you can hoist theFakeCamera` initialization out of the loop.

[lcc815]

Hi @njroussel,

Thanks for your quick reply. How can I hoist the FakeCamera initialization out of the loop? Like this:(?)

if __name__ == "__main__":
    sensor = mi.load_dict({
        'type': 'fakecamera',
        'to_world': mi.ScalarTransform4f.look_at(
                    origin=[0, 0., 1],
                    target=[0, 0, 0],
                    up=[0, -1, 0]
                ),
        'sampler': {
            'type': 'independent',
            'sample_count': 16
        },
        'film': {
            'type': 'hdrfilm',
            'width': 1920,
            'height': 1080,
            'pixel_format': 'rgb',
            'component_format': 'float32'
        }
    })
    for i in range(1000):
        scene = mi.load_file("scenes/cbox.xml")
        img = mi.render(scene, sensor=sensor, spp=16)  

But I want to update the to_world matrix every loop.

[njroussel]

You can edit any parameter of a plugin by using mi.traverse(). However, the plugin must explicitly define what its parameters are by overiding it's own traverse method. You can see an example of this in this tutorial: https://mitsuba.readthedocs.io/en/stable/src/others/custom_plugin.html

[lcc815]

Thanks for your advice.

I've tried to update the camera pose in my custom sensor class by modifying the to_world parameter. However, I noticed that when using mi.traverse() to inspect my custom sensor class, there is no to_world key available, while this key exists in the official-implemented perspective camera. I would like to know: Is there a specific way to register the to_world parameter in a custom sensor class?

Any help or suggestions would be greatly appreciated. Thank you!

UPDATE: According to here, I think the to_world parameter is not exposed. Does this means I have to initialize my custom sensor every time to update the to_world matrix (which will lead to memory leak)?

[lcc815]

Aha, seems you have replied in your last reply... Thanks, problem solved!

[njroussel]

Just to expand on what I said previously:

---

However, the plugin must explicitly define what its parameters are by overriding it's own traverse method.

That's the important bit. You need to override the method in your class definition.

In the tutorial I linked, the MyBSDF class overrides it as follows:

def traverse(self, callback):
        callback.put_parameter('tint', self.tint, mi.ParamFlags.Differentiable)