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?


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): = torch.randn(100, 3, 24, 24)
    def __getitem__(self, index):
        # Get current image
        image =[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 = torch.stack(patches)
        return patches
    def __len__(self):
        return len(

dataset = MyDataset()

loader = DataLoader(

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