Ways to clear memory.

[maxfrei750]

I'm currently fighting with a high memory consumption when using the path integrator together with transparent materials and the wrap_ad decorator (from drjit to torch).

Using the following code, called between iterations, I was able to significantly reduce the problem:

del image
gc.collect()
gc.collect()

torch.cuda.empty_cache()

With image being a previously rendered image. While looking for solutions, I found the clean_up() method in the code base:

mitsuba3/src/conftest.py

Lines 45 to 64 in 152352f

	def clean_up():
	'''
	Ensure no leftover instances from other tests are still in registry, wait
	for all running kernels to finish and reset the JitFlag to default. Also
	periodically frees the malloc cache to prevent the testcases from hogging
	all system memory.
	'''
	gc.collect()
	gc.collect()
	dr.kernel_history_clear()
	dr.flush_malloc_cache()
	dr.malloc_clear_statistics()
	dr.flush_kernel_cache()
	if hasattr(dr, 'sync_thread'):
	dr.sync_thread()
	dr.registry_clear()
	dr.set_flags(dr.JitFlag.Default)

Now I'm wondering if some of these methods might help to further alleviate my problem. Unfortunately, I couldn't find a documentation for them. Therefore, I'd be grateful, if somebody could shed some light on them and assess if they are suitable to reduce memory load while still being inside an optimization loop, i.e. not mess with the optimization.

Thanks!

[njroussel]

Hi

The path integrator has a significant memory/efficiency cost to it when used in AD settings. I'd strongly recommend using prb instead, take a look at figure 6 of the Path Replay Backpropagation paper.
The wrap_ad function has some caveats to it but they rather have to do with runtime than memory consumption usually.

I can't find the thread about it anymore, but typically torch also likes to reserve a huge chuck of your VRAM. So that might also be an avenue to look into.

The clean_up() method you found is very "strict". Because it's used between unit tests, it practically flushes the entire JIT's state which is rarely something you want. Anyway, this should not be a significant amount of memory and manually calling del variable;gc.collect();gc.collect() is the intended way to force variables to be deleted.

[maxfrei750]

Thanks for elaborating! I'd indeed prefer to use prb, but I'm facing a problem with exploding gradients (or something else that causes my parameters to sharply explode/drop), when I use it to optimize transparency, which is sort of expected, if I understand the documentation correctly:

the properties of ideal specular objects (e.g., the IOR of a glass vase) cannot be optimized.

I tried to experiment with roughness and the prb_reparam integrator. But I'm not sure, if I understood the implications of the remark in the docs correctly.

[njroussel]

You might want to take a look at the discussion over here or read this paper. To understand this notion of attached/detached sampling.

[njroussel]

Oh I hadn't yet seen your comment in this thread: #379
There's a lot going on in this context, I can only recommend reading a bit more of the literature that has been linked in the documentation and these discussions.

[maxfrei750]

Will do. Thanks for pointing me into the right direction.