Adding a caching dataloader for 3D dataset

[edyoshikun]

Summary

This PR adds a dataloader that caches the datasets in RAM and avoids using the SliceWindowDataset(). Here we cache the dataset on epoch 1. Ideally, the datasets fit within the 2TBs available per node. One can use btop to inspect the resource usage.

The z-slicing for the volumes is set as a parameter num_z_slices. We leave the slicing to the configurable MONAI transforms.

Testing

Working

• Tested the hcs_ram.CachedDataset
• Tested hcs_ram.CacheDataModule
• Tested running UNeXt2 with a small toy dataset locally (M1) and gpu-sm02 nodes.
• Running the model with dev_mode=True, which runs the model for 1 epoch.

Not Working

• Running UNeXt2 with the actual datasets (500GB+) using one GPU (H100), one node
• Running UNeXt2 with the actual datasets (500GB+) using 4 GPU (H100), one node