Custom BatchSampler for two-step mini-batch

Hi All,

In the data preparation phase for my network, I read an image one at a time, and then, I want to extract several patches from this image as my mini-batch. In other words, the data preparation consists of two steps: 1) read an image and 2) extract random patches to form the mini-match.

What’s the proper way to use BatchSampler to implement this?

Thanks,
Saeed

You could just sample in __getitem__ and stack the patches into the batch dimension:
Here is a small example sampling 5x5 patches. These patches are returned as a 4-dimensional tensor from __getitem__. During the training you could push these patches into the batch dimension with view.

class MyDataset(Dataset):
    def __init__(self):
        self.data = torch.randn(100, 3, 24, 24)
        
    def __getitem__(self, index):
        # Get current image
        image = self.data[index]
        # Sample patches
        patches = self.sample_pathes(image)
        
        return patches
    
    def sample_pathes(self, image):
        # Your sampling logic
        size = 5
        patches = []
        for i in range(5):
            patch = image[:, i:i+size, i:i+size]
            patches.append(patch)
        patches = torch.stack(patches)
        
        return patches
    
    def __len__(self):
        return len(self.data)


dataset = MyDataset()

loader = DataLoader(
    dataset,
    batch_size=10,
    shuffle=False,
    num_workers=2
)

loader_iter = iter(loader)
x = loader_iter.next()
x = x.view(-1, 3, 5, 5)
1 Like