Make sure batch_size is the same as defined in DataLoader, even after transforms

[simonebonato]

Hi all,
I'm working on a project with 3d head CT scans, and before feeding them to my segmentation model I want to perform some augmentations using monai transforms. The problem is that after I apply RandSpatialCropSamplesd() with num_samples=4, even if I set a batch_size of 1 in the DataLoader I still get an effective batch_size of 4, and since the data is too big I cannot load it onto the GPU.

Is there a way to overcome this issue? I tried changing the collate_fn and using other transforms but without success.

The transforms I am using are the following:

Compose(
                    [
                        LoadImaged(keys=["image", "label"]),
                        EnsureChannelFirstd(keys=["image", "label"]),
                        Orientationd(keys=["image", "label"], axcodes="RAS"),
                        ScaleIntensityRanged(keys=["image"], a_min=-175, a_max=250, b_min=0.0, b_max=1.0, clip=True),
                        CastToTyped(keys=["label"], dtype=int),
                        OneOf(
                            transforms = [
                                CropForegroundd(keys=["image", "label"], 
                                                source_key="label",
                                                select_fn=lambda x: x > 0.5,
                                                margin=(50, 50, 50)
                                                ),
                                RandSpatialCropSamplesd(keys=["image", "label"], 
                                                                roi_size=self.img_size,
                                                                random_size=False,
                                                                num_samples=1,),
                            ],
                            weights=[0., 1.]
                        ),
                        ResizeWithPadOrCropd(keys=["image", "label"], 
                                            spatial_size=self.img_size
                                     ),
                        RandFlipd(
                            keys=["image", "label"],
                            spatial_axis=[0],
                            prob=0.10,
                        ),
                        RandFlipd(
                            keys=["image", "label"],
                            spatial_axis=[1],
                            prob=0.10,
                        ),
                        RandFlipd(
                            keys=["image", "label"],
                            spatial_axis=[2],
                            prob=0.10,
                        ),
                        RandRotate90d(
                            keys=["image", "label"],
                            prob=0.10,
                            max_k=3,
                        ),
                        RandShiftIntensityd(
                            keys=["image"],
                            offsets=20,
                            prob=0.50,
                        ),
                        ToTensord(keys=["image", "label"]),
                    ]
)

[wyli]

perhaps you can try wrapping the multi-sample dataset with a monai.data.PatchDataset:

dataset = monai.data.Dataset(filenames, transform=xform)  # xform is the Compose with RandSpatialCropSamplesd
p_dataset = monai.data.PatchDataset(dataset, patch_func=lambda x: x, samples_per_image=4)
loader = monai.data.DataLoader(p_dataset, num_workers=1, batch_size=3, shuffle=False)

the output batch size would be 3 in this case. (I haven't tested this code in an end-to-end workflow, please double-check the details such as randomness, multiprocessing...)

for more info:

MONAI/monai/data/grid_dataset.py

Lines 202 to 226 in 8eceabf

	class PatchDataset(Dataset):
	"""
	returns a patch from an image dataset.
	The patches are generated by a user-specified callable `patch_func`,
	and are optionally post-processed by `transform`.
	For example, to generate random patch samples from an image dataset:
	.. code-block:: python
	import numpy as np
	from monai.data import PatchDataset, DataLoader
	from monai.transforms import RandSpatialCropSamples, RandShiftIntensity
	# image dataset
	images = [np.arange(16, dtype=float).reshape(1, 4, 4),
	np.arange(16, dtype=float).reshape(1, 4, 4)]
	# image patch sampler
	n_samples = 5
	sampler = RandSpatialCropSamples(roi_size=(3, 3), num_samples=n_samples,
	random_center=True, random_size=False)
	# patch-level intensity shifts
	patch_intensity = RandShiftIntensity(offsets=1.0, prob=1.0)
	# construct the patch dataset
	ds = PatchDataset(dataset=images,

[simonebonato]

Thanks for the reply!

So what you suggested seems to work, but only in the case RandSpatialCropSamplesd with num_samples=4 is used; but as you can see in the the transform I use I wanted to put

OneOf(
                            transforms = [
                                CropForegroundd(keys=["image", "label"], 
                                                source_key="label",
                                                select_fn=lambda x: x > 0.5,
                                                margin=(50, 50, 50)
                                                ),
                                RandSpatialCropSamplesd(keys=["image", "label"], 
                                                                roi_size=self.img_size,
                                                                random_size=False,
                                                                num_samples=1,),
                            ],
                            weights=[0., 1.]
                        ) 

I set the weight [0., 1.] just for testing out the code, but in a real case scenarion I would put something like [0.8, 0.2], hence when the CropForegroundd transform is used, if I use your solution with the PatchDataset I get the following error:

RuntimeWarning: Caught RuntimeWarning in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/home/student/.conda/envs/3d_seg/lib/python3.9/site-packages/torch/utils/data/_utils/worker.py", line 302, in _worker_loop
    data = fetcher.fetch(index)
  File "/home/student/.conda/envs/3d_seg/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 58, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/student/.conda/envs/3d_seg/lib/python3.9/site-packages/torch/utils/data/_utils/fetch.py", line 58, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/home/student/.conda/envs/3d_seg/lib/python3.9/site-packages/monai/data/dataset.py", line 107, in __getitem__
    return self._transform(index)
  File "/home/student/.conda/envs/3d_seg/lib/python3.9/site-packages/monai/data/grid_dataset.py", line 272, in _transform
    raise RuntimeWarning(
RuntimeWarning: `patch_func` must return a sequence of length: samples_per_image=4.

Since the output of the transforms will have 1 sample only and not 4...

[wyli]

yeah, the PatchDataset requires a fixed samples_per_image in order to compute the index

MONAI/monai/data/grid_dataset.py

Lines 264 to 265 in 8eceabf

	def __len__(self) -> int:
	return len(self.data) * self.samples_per_image

for dynamic number of samples per image, you may look into PatchIterd

MONAI/monai/data/grid_dataset.py

Lines 81 to 90 in 8eceabf

	class PatchIterd:
	"""
	Dictionary-based wrapper of :py:class:`monai.data.PatchIter`.
	Return a patch generator for dictionary data and the coordinate, Typically used
	with :py:class:`monai.data.GridPatchDataset`.
	Suppose all the expected fields specified by `keys` have same shape.
	Args:
	keys: keys of the corresponding items to iterate patches.
	patch_size: size of patches to generate slices for, 0/None selects whole dimension