homogenous_batching.py

class CustomBatchSampler(Sampler):
    def __init__(self, batch_size, dataset):
        self.sampler = SequentialSampler(dataset)
        self.batch_size = batch_size
        self.drop_last = False
        self.dataset = dataset
    def __iter__(self):
        batch = []
        for idx in self.sampler:
            batch.append(idx)
            if len(batch) == self.batch_size:
                yield batch
                batch = []
            current_class = self.dataset[idx][-1]
            try:
                next_class = self.dataset[idx+1][-1]
            except IndexError:
                next_class = current_class
            
            if batch and next_class != current_class:
                yield batch
                batch = []
    def __len__(self):
        if self.drop_last:
            return len(self.sampler) // self.batch_size
        else:
            return (len(self.sampler) + self.batch_size - 1) // self.batch_size
          
 
# Use like this:
dataset = ImageFolder(...)
sampler = CustomBatchSampler(batch_size=16, dataset=dataset)
loader = DataLoader(dataset, batch_sampler=sampler, num_workers=4)