Dependent Batches in DataSet
Category: vision

melgor: What do you think about this strategy?

apaszke: Hey, sorry for a late reply. I didn’t really understand your solution, sampler is not a function, but an iterable (you can think of it as a stream of indices). I think that one way to impement what you want would be to:

melgor: Thanks for the reply. Here is my code: <SCODE>from PIL import Image

import sys
import collections

def collate_cat(batch):
    '''concatenate sub-batches along first dimension'''
    if torch.is_tensor(batch[0]):
        return torch.cat(batch,0)
    elif isinstance(batch[0], collections.Iterable):
        # if each batch element is not a tensor, then it should be a tuple
        # of tensors; in that case we collate each element in the tuple
        transposed = zip(*batch)
        return [collate_cat(samples) for samples in transposed]

    raise TypeError(("batch must contain tensors, numbers, or lists; found {}"
                     .format(type(batch[0]))))

class RandomSampler(object):
    """Samples num_iters times the idxes from class
    Arguments:
        data_source (Dataset): dataset to sample from
    """

    def __init__(self, data_source, num_iters):
        self.num_samples = len(data_source)
        self.num_iters   = num_iters

    def __iter__(self):
        return iter(torch.cat([torch.randperm(self.num_samples) for i in range(self.num_iters)], 0).long())

    def __len__(self):
        return self.num_iters
    
    
class TripletDataset(datasets.ImageFolder):
    def __init__(self, root, transform=None, target_transform=None, num_iters = 1000, \
                     class_per_batch = 10, example_per_class = 6):
        super(self.__class__, self).__init__(root, transform, target_transform)
        self.class_per_batch   = class_per_batch
        self.example_per_class = example_per_class
        
        # Create a dictionary which will contain idxes of each class separately\
        self.class_dict = defaultdict(list)
        for idx, (path, target) in enumerate(self.imgs):
            self.class_dict[target].append(idx)
               
        
    def __getitem__(self, index):
        class_data  = self.class_dict[index]
        num_example = min(len(class_data), self.example_per_class)
        shuffle     = torch.randperm(len(class_data))
        idxes       = shuffle[:num_example]
        list_imgs    = list()
        list_targets = list()
        for idx in idxes:
            path, target = self.imgs[class_data[idx]]
            img = self.loader(os.path.join(self.root, path))
            if self.transform is not None:
                img = self.transform(img)
            if self.target_transform is not None:
                target = self.target_transform(target)
            list_imgs.append(img)
            list_targets.append(target)
            
        return torch.stack(list_imgs, 0), torch.LongTensor(list_targets) #img, target

    def __len__(self):
        return len(self.class_dict.keys())



data = TripletDataset("MNIST/train/",transforms.Compose([transforms.ToTensor()]))
train_loader = torch.utils.data.DataLoader(
        data,
        batch_size=3, shuffle=True,
        sampler = RandomSampler(data, 100 * 3),
        num_workers=0, pin_memory=True,
        collate_fn = collate_cat)

for img, target in train_loader:
    print (img.size(), target.size())
<ECODE>

apaszke: About the iterator, it would be slightly more efficient to recreate the permutations only once they’re needed, not when the iterator is instantiated, but as long as your dataset isn’t huge it probably doesn’t matter too much.