火炬 - 对每个数据集进行不同的子采样并将它们连接起来

Question

I have two datasets, but one is larger than the other and I want to subsample it (resample in each epoch).

I probably cannot use dataloader argument sampler, as I would pass to Dataloader the already concatenated dataset.

How do I achieve this simply?

I think one solution would be to write a class SubsampledDataset(IterableDataset) which would resample every time __iter__ is called (each epoch).

(Or better use a map-style dataset, but is there a hook that gets called every epoch, like __iter__ gets?)

Answer 1

You can stop to iterate by raising StopIteration . This error is caught by Dataloader and simply stop the iteration. So you can do something like that:

class SubDataset(Dataset):
    """SubDataset class."""
    def __init__(self, dataset, length):
        self.dataset = dataset
        self.elem = 0
        self.length = length

    def __getitem__(self, index):
        self.elem += 1
        if self.elem > self.length:
            raise StopIteration  # caught by DataLoader
        return self.dataset[index]

    def __len__(self):
        return len(self.dataset)


if __name__ == '__main__':
    torch.manual_seed(0)
    dataloader = DataLoader(SubDataset(torch.arange(10), 5), shuffle=True)
    for x in dataloader:
        print(x)  # 6, 7, 1, 4, 2
    print(len(dataloader))  # 10!!

Note that setting __len__ to self.length will cause a problem because dataloader will use only indices between 0 and length-1 (that is not what you want). Unfortunately I found nothing to set the actually length without having this behaviour (due to Dataloader restriction). Thus be careful: len(dataset) is the original length and dataset.length is the new length.

Answer 2

This is what I have so far (untested). Usage:

dataset1: Any = ...
# subsample original_dataset2, so that it is equally large in each epoch
dataset2 = RandomSampledDataset(original_dataset2, num_samples=len(dataset1))

concat_dataset = ConcatDataset([dataset1, dataset2])

data_loader = torch.utils.data.DataLoader(
    concat_dataset,
    sampler=RandomSamplerWithNewEpochHook(dataset2.new_epoch_hook, concat_dataset)
)

The result is that the concat_dataset will be shuffled each epoch (RandomSampler), in addition, the dataset2 component is a new sample of the (possibly larger) original_dataset2, different in each epoch.

You can add more datasets to be subsampled by doing instead of:

sampler=RandomSamplerWithNewEpochHook(dataset2.new_epoch_hook

this:

sampler=RandomSamplerWithNewEpochHook(lambda: dataset2.new_epoch_hook and dataset3.new_epoch_hook and dataset4.new_epoch_hook, ...

Code:

class RandomSamplerWithNewEpochHook(RandomSampler):
    """ Wraps torch.RandomSampler and calls supplied new_epoch_hook before each epoch. """
    
    def __init__(self, new_epoch_hook: Callable, data_source: Sized, replacement: bool = False,
                 num_samples: Optional[int] = None, generator=None):
        super().__init__(data_source, replacement, num_samples, generator)
        self.new_epoch_hook = new_epoch_hook

    def __iter__(self):
        self.new_epoch_hook()
        return super().__iter__()


class RandomSampledDataset(Dataset):
    """ Subsamples a dataset. The sample is different in each epoch.

    This helps when concatenating datasets, as the subsampling rate can be different for each dataset.
    
    Call new_epoch_hook before each epoch. (This can be done using e.g. RandomSamplerWithNewEpochHook.)

    This would be arguably harder to achieve with a concatenated dataset and a sampler argument to Dataloader. The
    sampler would have to be aware of the indices of subdatasets' items in the concatenated dataset, of the subsampling 
    for each subdataset."""
    def __init__(self, dataset, num_samples, transform=lambda im: im):
        self.dataset = dataset
        self.transform = transform
        self.num_samples = num_samples

        self.sampler = RandomSampler(dataset, num_samples=num_samples)
        self.current_epoch_samples = None

    def new_epoch_hook(self):
        self.current_epoch_samples = torch.tensor(iter(self.sampler), dtype=torch.int)

    def __len__(self):
        return self.num_samples

    def __getitem__(self, item):
        if item < 0 or item >= len(self):
            raise IndexError

        img = self.dataset[self.current_epoch_samples[item].item()]

        return self.transform(img)