Cache location used by CacheDataset

[dianemarquette]

Hi!

I originally trained my neural networks using a GeForce RTX 2080 Super (with 8Gb VRAM) and 32 GB RAM. Recently, I changed setups, and I now have a GeForce RTX 3090 (with 24Gb VRAM) but only 16 GB RAM.

I tried to train a neural network using MONAI's CacheDataset class on my new setup (with the same code, same conda environment, same dataset). It worked fine on my old setup, but I got an OOM error in the middle of my training set loading process on the new setup.

Until now, I had assumed that CacheDataset used the GPU VRAM to cache the transformed scans.

Given the bug, I wonder whether the limited amount of RAM causes this OOM on my new setup.

My questions:

• Does CacheDataset use the PC RAM (instead of the GPU VRAM)?
• If so, how can I make CacheDataset take advantage of my GPU VRAM (which is largely unused when the OOM occurs)?

Thanks in advance for your insights :)

[Nic-Ma]

Hi @dianemarquette ,

Thanks for your interest here and experiments with MONAI.

1. It depends on your transforms where to cache the data, if you used ToDevice("cuda:0") it will cache the PyTorch Tensor to GPU RAM, otherwise, in PC RAM like your case, it's just numpy arrays.
2. Please refer to this tutorial on how to cache in GPU RAM and execute following transforms on GPU directly for better perf:
   https://github.com/Project-MONAI/tutorials/blob/master/acceleration/fast_training_tutorial.ipynb
   Please note that it's still under development, not all the MONAI transforms support GPU Tensor operations.

Thanks.

[dianemarquette]

Thanks @Nic-Ma for the clear answer!