Memory behavior of PersistentDataset

[bazorb]

I am working on a modified Radiology App of MONAI Label and a modified segmentation model. I'm frequently experiencing out of memory errors (x4 16GB ram GPUs). I found the best luck with PersistentDataset setting in the 3DSlicer plugin. CacheDataset and SmartCacheDataset will not run. I'm using N=25 CTs as my training set and a batch size for training and validation of N=1.

To help debug these memory issues, I made a custom transform Memoryd to report the GPU memory utilization (a kind of poor man's profiler based on gpu_memory_map). I notice that as the training and validation proceed the memory utilization steadily rises. The memory hits a maximum at Epoch 2 or 3 of the Ignite engine. When I use distributed training, the memory utilization for each GPU is higher and get an OOM error. This means I either reduce my dataset size or use one GPU.

My question is what is the memory handling behavior of PersistentDataset? My expectation from the documentation is that the non-randomizable transforms (up to EnsureTyped below) will be saved to a file location and cached on disc. Subsequent epochs and transforms will pick up from the disc and carry out the following transforms on the GPU. I expect the GPU memory to reset after each batch is completed and allow plenty of memory for each batch. Instead the PersistentDataset is acting more like how I expect CacheDataset to behave where data is stored on the GPU memory for the entire train task.

Any advice here? Thank you in advance.

def train_pre_transforms(self, context: Context):
        return [
            LoadImaged(keys=("image", "label")),
            EnsureChannelFirstd(keys=("image", "label")),
            Spacingd(
                keys=("image", "label"),
                pixdim=self.target_spacing,
                mode=("bilinear", "nearest"),
            ),
            ScaleIntensityRanged(keys="image", a_min=-57, a_max=164, b_min=0.0, b_max=1.0, clip=True),
            CropForegroundd(keys=("image", "label"), source_key="image"),
            EnsureTyped(keys=("image", "label"), device=context.device),
            RandCropByPosNegLabeld(
                keys=("image", "label"),
                label_key="label",
                spatial_size=self.roi_size,
                pos=1,
                neg=1,
                num_samples=4,
                image_key="image",
                image_threshold=0,
            ),
            RandShiftIntensityd(keys="image", offsets=0.1, prob=0.5),
            SelectItemsd(keys=("image", "label")),
            Memoryd(keys=("image")),
        ]

[KumoLiu]

Hi @bazorb, I think your understanding is right, maybe you could share a short piece of code to reproduce the problem you are pointing out and we can analyze it further.
Thanks!

[bazorb]

@KumoLiu thank you for the reply.

I've attached a sample segment of code. It is only for 1 GPU as the whole MONAI Label infrastructure is too complex to replicate simply. It includes the Memoryd transform (although there is some bug of repeated prints). Using this code a did several experiments to monitor how memory usage changes with different options. Seems like with the synthetic data there is about 2 GB memory usage regardless of much else. Even doubling the image size in each dimension did not significantly impact the memory consumption.

At this stage my question is, when is memory allocated to the GPU and when is it released? What is needed to ensure that the GPU will not cache multiple samples at once which I believe is the origin of my memory issues.

sample_trainer.zip

[KumoLiu]

Hi @bazorb, the data is allocated to the GPU after EnsureTyped(keys=("image", "label"), device=device) this transform. PersistentDataset will save the results after all non-random transforms to disk. From your code I noticed that you used EnsureTyped to allocated data to the GPU, actually if you are using PersistentDataset, the data should no longer be loaded onto the GPU using this transform.
You could refer to this tutorial and find more details about these two dataset.

Hope it can help you, thanks!

[Nic-Ma]

Hi @bazorb ,

GPU usually has some memory pool mechanism that may not release a memory buffer immediately unless it will be OOM soon. So what your Memoryd reported may not be very accurate.
I think your real issue is why distributed data parallel uses more memory on each GPU? Please share your test program then @KumoLiu can help analyze further.

Thanks.