Project-MONAI / Project-MONAI/MONAI

Memory leak in monai.transforms.distance_transform_edt

Open
#7,727 1 comment 1 reaction 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
8.7k
Forks
1.6k
Avg merge
5d 1h
Merged PRs (30d)
20

Description

Describe the bug
There seems to be a memory leak in monai.transforms.distance_transform_edt when cuCIM is not available.

To Reproduce
cuCIM must not be available to reproduce

import torch
from monai.transforms import distance_transform_edt

nb_interations = 2000
for i in range(nb_interations):
    mask = torch.randint(low=0, high=2, size=(10, 10, 10,), device='cuda')
    distance_transform_edt(mask)

print(torch.cuda.memory_summary(device=None, abbreviated=True))

Expected behavior
The GPU memory footprint increases with nb_iterations. This should not happen. This does not happen when cuCIM is installed.

Additional context
The problem seems to be an exception leak in optional_import:

https://github.com/Project-MONAI/MONAI/blob/56508992c32e328667740e18f99b82a6a19f2c2a/monai/utils/module.py#L398-L408

where tb does not get cleared. tb holds references that prevent the mask variable from being deleted. I think this could be a problem with other optional_import calls, too.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with monai/utils/module.py around lines 398-408 and inspect optional_import, then reproduce the leak through monai.transforms.distance_transform_edt with cuCIM unavailable using the provided CUDA loop. Done means the repeated calls no longer cause GPU memory to grow, including for the reported fallback path.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
performance
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Clearly specified
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.