LabelSampler Memory Consumption

Hi All, I am using Label Sampler as follows:

psize = (144,144,80)
sampler = torchio.data.LabelSampler(psize)
patches_queue = torchio.Queue(subjects_dataset, max_length=1,
                                samples_per_volume=1,
                                sampler=sampler, num_workers=4,
                                shuffle_subjects=False, shuffle_patches=True, verbose=False)

The LabelSampler seems to consume CPU memory upwards of 120GB. I have 160ish training images. I am using about 6 augmentations with probability = 0.1.

However, when I use the exact same configuration, but with UniformSampler, the CPU memory requirement is much lesser.

Any reason this is happening? Thanks for your help, Megh

[meghbhalerao]

Thanks a lot @fepegar for your detailed analysis.

I observed one thing. When I am doing padding for LabelSampler and WeightedSampler as per our discussion here: https://github.com/fepegar/torchio/issues/392 - like this:

psize = (144,144,80)
psize_pad = list(np.asarray(np.round(np.divide(psize,2)), dtype=int))
padder = Pad(psize_pad, padding_mode = 'symmetric') # for modes: https://numpy.org/doc/stable/reference/generated/numpy.pad.html
subject = padder(subject)

By doing this padding I observed that the memory requirement shoots up by a lot (>125GB). I am trying to investigate what the issue might be. Thanks again, Megh

[meghbhalerao]

Thank you @fepegar. Sure, I will try using WeightedSampler and let you know. Will try making a dummy dataset and let you know. Thanks again