Memory leak with `h5py` from `pip` and conversion to `torch.Tensor`

[Breeze-Zero]

I recently tried to do some experiments on my model with multi-coil FastMRI brain data. Due to the need for flexibility (and also because I don't have the extra time to learn how to use Pytorch lighting), I didn't use Pytorch Lighting directly. Instead, I chose normal Pytorch, but during the iterating process, I only set num_worker=2, and my memory footprint was quite large at the beginning. As the number of iterations increased, an error occurred:
RuntimeError: DataLoader worker (PID 522908) is killed by signal: killed.
I checked the training codes of other parts, but no obvious memory accumulation error was found. Therefore, I thought there was a large probability of a problem in siliceDataset. I simply used "pass" to traverse the Dataloader loop, and found that the memory occupation kept rising.

[mmuckley]

Hello @834799106, thanks for putting an issue here.

Based on your error I doubt the SliceDataset class is the issue. For one, your program is not being terminated due to memory, it is being terminated because some process killed the overall program. Also, if you look at the __getitem__ function, you can see that there are no side-effects. Everything the function creates should be returned to the calling function or destroyed.

In order to verify a memory leak we will need you to give us a reproducible example for your case since you're not using the PyTorch Lightning modules. Also, please let us know what version of PyTorch you are using and any information you have on the memory usage throughout an epoch. Note: high memory at the start might be expected, as you have your model in memory. There is also some metadata about the dataset that is precomputed and stored in memory.

[soumickmj]

Hi @mmuckley I was about to file an issue for a memory leak. I'm not sure about the issue of @834799106 though.
I have created a small pieace of code for reproducing.

`
from fastmri.data.transforms import UnetDataTransform
from fastmri.data import SliceDataset
from fastmri.data.subsample import create_mask_for_mask_type
import os
from torch.utils.data import DataLoader
from tqdm import tqdm

mask_func = create_mask_for_mask_type(
mask_type_str="random", center_fractions=[0.08], accelerations=[8])

root_gt = "/data/project/fastMRI/Brain/multicoil_train"

sd = SliceDataset(root=root_gt, challenge="multicoil",
transform=UnetDataTransform("multicoil", mask_func=mask_func, use_seed=False), use_dataset_cache=True,
dataset_cache_file=f"{os.path.dirname(root_gt)}/dataset_cache_{os.path.basename(root_gt)}.pkl")
dl = DataLoader(sd, batch_size=1, shuffle=False, num_workers=10)

for e in tqdm(dl):
del e
pass
`
I'm currently using the latest git pull of fastMRI.

While running this code, I was monitoring the memory usage, even though I'm deleting the varible, it is still increasing the memory usage constantly. Originally, this was part of my other pipeline where I'm only using the SliceDataset and not the whole Lightning module. If you would like to have a look, this is the code: https://github.com/soumickmj/NCC1701/blob/main/Engineering/datasets/fastMRI.py

I was originally thinking maybe my code is creating the leak. But with other dataset mode (different code for reading other datasets) of the same NCC1701 pipeline of mine did not create the leak.
Then I wrote that small script to try to see if the leak is still there when my pipeline is not involved.

[soumickmj]

I also got a similar behaviour while using the Data Module.

`
from fastmri.data.transforms import UnetDataTransform
from fastmri.pl_modules import FastMriDataModule
from fastmri.data import SliceDataset
from fastmri.data.subsample import create_mask_for_mask_type
import os
from torch.utils.data import DataLoader
from tqdm import tqdm

mask_func = create_mask_for_mask_type(
mask_type_str="random", center_fractions=[0.08], accelerations=[8])

root_gt = "/data/project/fastMRI/Brain"

data_module = FastMriDataModule(data_path=root_gt,
challenge="multicoil",
train_transform=UnetDataTransform("multicoil", mask_func=mask_func, use_seed=False),
val_transform=UnetDataTransform("multicoil", mask_func=mask_func, use_seed=False),
test_transform=UnetDataTransform("multicoil", mask_func=mask_func, use_seed=False),
batch_size=1,
num_workers=10)
dl = data_module.train_dataloader()

for e in tqdm(dl):
del e
pass

`

[mmuckley]

Hello @soumickmj, I ran your script on the knee validation data with memory-profiler and memory usage peaked pretty early a little bit less than 5 GB (see attached), staying flat for the rest of the entire dataset afterwards (which does not suggest a leak).

Perhaps you could try running on your system to verify with PyTorch 1.10?

This is the code:

from fastmri.data.transforms import UnetDataTransform
from fastmri.data import SliceDataset
from fastmri.data.subsample import create_mask_for_mask_type
import os
from torch.utils.data import DataLoader
from tqdm import tqdm


@profile
def main():
    val_path = "/path/multicoil_val"
    mask_func = create_mask_for_mask_type(
        mask_type_str="random", center_fractions=[0.08], accelerations=[8]
    )

    sd = SliceDataset(
        root=val_path,
        challenge="multicoil",
        transform=UnetDataTransform("multicoil", mask_func=mask_func, use_seed=False),
    )
    dl = DataLoader(sd, batch_size=1, shuffle=False, num_workers=10)

    for e in tqdm(dl):
        del e
        pass


if __name__ == "__main__":
    main()

You can run with mprof run --include-children file.py.

[Breeze-Zero]

hi @mmuckley，I copied the above code to my machine, but modified batch Szie. The figure below is the result of mprof run --include-children file.py

I didn't even finish an epoch before broke off。

This is the code:

from fastmri.data.transforms import UnetDataTransform
from fastmri.data import SliceDataset
from fastmri.data.subsample import create_mask_for_mask_type
import os
from torch.utils.data import DataLoader
from tqdm import tqdm


@profile
def main():
    val_path = "/data2/fastmri/mnt/multicoil_val"
    mask_func = create_mask_for_mask_type(
        mask_type_str="random", center_fractions=[0.08], accelerations=[8]
    )

    sd = SliceDataset(
        root=val_path,
        challenge="multicoil",
        transform=UnetDataTransform("multicoil", mask_func=mask_func, use_seed=False),
    )
    dl = DataLoader(sd, batch_size=4, shuffle=False, num_workers=10)

    for e in tqdm(sd):
        del e
        pass


if __name__ == "__main__":
    main()

Maybe it's PyTorch version that's causing the problem. My PyTorch version is 1.8.1+ cu111

[soumickmj]

Sorry @mmuckley I also got the same problem after running memory profiler.
I used to different versions of PyTorch.
On the contrary to @834799106 I am using more latest versions of PyTorch.

With PyTorch 1.10.2 py3.9_cuda11.3_cudnn8.2.0_0 I got:-

With PyTorch 1.11.0.dev20220129 py3.9_cuda11.3_cudnn8.2.0_0 (pytorch-nightly), due to the features I use, I usually need this version for my work:-

I did not run till the very end as it was continously increasing and would have crashed the server again which has 250GB of RAM. So I don't feel its related to the PyTorch version.

Just to let you know: the OS is Ubuntu 20.04.3 LTS and the python version is 3.9.7

[soumickmj]

sd = SliceDataset(
        root=val_path,
        challenge="multicoil",
        transform=UnetDataTransform("multicoil", mask_func=mask_func, use_seed=False),
    )
    dl = DataLoader(sd, batch_size=1, shuffle=False, num_workers=10)

    for e in tqdm(sd):
        del e
        pass

Hi @mmuckley
In your code, I just noticed that you are looping over the dataset directly without the dataloader, whereas in my code I'm running over DataLoader. Can you please also try that out too?
In my case, both shown similar behaviour. I also tested with 0,1,3 as number of workers - all the same.
I created another conda env with python 3.7.11 and torch 1.8.2, got similar behaviour as well.

Am I doing something wrong somewhere?

[mmuckley]

Hello @soumickmj, I copied the wrong code. The paste I showed was with the dataloader, not the dataset. This is what I get with dataset. You can see the max is about 720 MB.

[soumickmj]

Hello @soumickmj, I copied the wrong code. The paste I showed was with the dataloader, not the dataset. This is what I get with dataset. You can see the max is about 720 MB.

Ah okay, no problem!
But still, in my case, I'm getting this constant increase in terms of memory usage as you can see from the plots.
Any suggestions?

[mmuckley]

One thing I notice is that you are both using Python 3.9. I could try Python 3.9 and check with that perhaps.

EDIT: Sorry, I see you also tried Python 3.7. Not sure what to do then...

[mmuckley]

Okay, I tested on Python 3.9 and I'm still getting the same behavior with batch_size=4.

To help a bit more with this I'm including my complete conda environments that I used for my tests. I'm using a custom Linux distribution based on 5.4.0-81-generic x86_64 on with an Intel Xeon E5-2698. @soumickmj @834799106, if either of you could try one of my conda environments maybe we could figure out if it's one of the packages.

Python 3.8 environment
Python 3.9 environment

[soumickmj]

Thanks @mmuckley
I have tested with your 3.9 conda env and it worked without a problem.
Your yml was missing fastMRI:
So I installed it using pip install git+https://github.com/facebookresearch/fastMRI.git

Then I created a "bare minimum" env without using your env, and this resulted in the old issue again.
For this env, I just installed pytorch 1.10.2 with cuda 11.3 and then directly fastMRI repo using git similar to the other one.
This conda environment doesn't contain anything which might not be requred.

I compared the versions. Initially, the version of numpy was different (1.21.2). So I switched to the one you have 1.20.3
Apart from all the extra packages which are there in your env, I couldn't see any difference.
Here is the yml file, zipped as it wasn't allowing me to upload yml here.
fMMem.zip

Do you have any idea what might be the reason?
Do I need some additional package that it is not complaining about but still required and causing this issue?

[Breeze-Zero]

I was in a similar situation. The problem was solved by installing your Py3.9 environment directly with Conda, however, creating a Py3.9 environment from Conda normally and pip install git+https://github.com/facebookresearch/fastMRI.git and pandas (It's not in the FastMRI package), the problem remains.

[mmuckley]

So my install process is as follows in a few bash commands:

conda create -n memory_test_py39 python=3.9
conda activate memory_test_py39
conda install anaconda
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
pip install -e .

Where the pip install -e . is in the fastMRI folder.

In that case I will try to reproduce now with your minimal environments.