Enforcing a batch size despite using RandCropByPosNegLabeld's num_samples

[jmlipman]

Hi,

I have very large images (say, 3000x3000x1000) and I want to extract several patches from each image for training. I was doing this with RandCropByPosNegLabeld by setting the parameter num_samples. However, this parameter dominates over the batch size. In other words, if I set my batch_size in the DataLoader to 2 and I extract 100 patches via RandCropByPosNegLabeld's num_samples=100, during the training loop it will try to give me a tensor of size [100xCxHxWxD]. Below I put a minimum working example illustrating this issue.

Is there any way in which I can extract a lot of patches, keep them in memory after being augmented, and use them iteratively during training?

Thanks!

import nibabel as nib
import numpy as np
from monai.transforms import (
        Compose,
        LoadImaged,
        RandCropByPosNegLabeld,
        )
from monai.data import DataLoader, CacheDataset

BATCH_SIZE = 2
NUM_SAMPLES = 10

# Create empty image
nib.save(nib.Nifti1Image(np.zeros((1, 200, 200, 200)), np.eye(4)), 'image.nii.gz')
trainFiles = [{'image': 'image.nii.gz', 'label': 'image.nii.gz'}]

transforms = Compose([
    LoadImaged(keys=['image', 'label']),
    RandCropByPosNegLabeld(keys=['image', 'label'],
                           spatial_size=[32, 32, 32],
                           label_key='label',
                           num_samples=NUM_SAMPLES)
    ])

tr_data = CacheDataset(data=trainFiles, transform=transforms, cache_rate=1.0)
tr_loader = DataLoader(tr_data, batch_size=BATCH_SIZE)

for tr_i, batch_data in enumerate(tr_loader):
    X, Y = batch_data['image'], batch_data['label']
    print(X.shape) 

With BATCH_SIZE = 2, NUM_SAMPLES = 10
Output:
torch.Size([10, 1, 32, 32, 32])

Expected/desired output:
torch.Size([2, 1, 32, 32, 32])
torch.Size([2, 1, 32, 32, 32])
torch.Size([2, 1, 32, 32, 32])
torch.Size([2, 1, 32, 32, 32])
torch.Size([2, 1, 32, 32, 32])

[KumoLiu]

Hi @jmlipman, since random transforms always generate different patches, if your epoch num is large, I think you've achieved what you want with the RandCropByPosNegLabel transform without the need to save the random patches.
Thanks!

[jmlipman]

Hi @KumoLiu Thanks for the answer. The problem is that the images are very large and I have few images (reading one image takes about 30 seconds), that's why I prefer to extract a lot of patches from the images. As a result, I may train for a few epochs but for a large number of iterations.

In the worst case, I could do a for loop inside the training loop to divide the large tensor from all the patches into mini-batches, but maybe there is a cleverer way to do this.

[KumoLiu]

Hi @jmlipman, got it, maybe you could try to use PatchDataset. Thanks!
https://github.com/Project-MONAI/MONAI/blob/dev/monai/data/grid_dataset.py#L202