Code Reproducibility

[JingTongsh]

Hello, I am trying to train your model with the MSD dataset, as you suggested.

I have encountered several problems during the process:

1. Dataset conversion. When converting the dataset.json to your version with a "prompt" for each subject, I wonder which prompt to use when there are more than one targets in the image. For instance, MSD Task01 BrainTumour has 3 classes of interest: edema, enhancing tumour and non-enhancing tumor. I guess your model works with one target at a time. So my question is:

Do you write all targets (or those existing in the current image) into the prompt?

Or do you copy the image and separate the ground truth mask with all classes to binary masks for each class? I mean, convert a pair {image, label} with $label\in {0,1,2,3}^{H\times W}$ to, say, 3 pairs

{image, label_1},
{image, label_2},
{image, label_3}

with each $label_k\in {0,1}^{H\times W}, k\in{1,2,3}$? I suppose this is the case, implemented a conversion script, and found issues below.

2. Data transform in loaders. When loading data from MSD datasets, there are errors like array shape mismatch, ROI larger than image, etc. I have fixed such errors for MSD Task01 BrainTumour, but found more errors in MSD Task05 Prostate. I'm not sure if the problem comes from my data conversion script. Hope you can provide a solution.

3. CUDA usage I wonder how much CUDA memory you used to train the model with automatic pathway modules from pretrained SwinUNETR. For MSD Task01 BrainTumour, I chose SwinTiny as the backbone, froze all its layers except input patch_embed, set batch_size=2 and my GPU with 48GiB went out of memory. It starts to work when I set batch_size=1, still occupying 33GiB.