| |
| |
| |
| |
|
|
| import logging |
|
|
| import numpy as np |
|
|
| from . import BaseWrapperDataset |
|
|
|
|
| logger = logging.getLogger(__name__) |
|
|
|
|
| class SubsampleDataset(BaseWrapperDataset): |
| """Subsamples a given dataset by a specified ratio. Subsampling is done on the number of examples |
| |
| Args: |
| dataset (~torch.utils.data.Dataset): dataset to subsample |
| size_ratio(float): the ratio to subsample to. must be between 0 and 1 (exclusive) |
| """ |
|
|
| def __init__(self, dataset, size_ratio, shuffle=False): |
| super().__init__(dataset) |
| assert size_ratio < 1 |
| self.actual_size = np.ceil(len(dataset) * size_ratio).astype(int) |
| self.indices = np.random.choice( |
| list(range(len(self.dataset))), self.actual_size, replace=False |
| ) |
| self.shuffle = shuffle |
| logger.info( |
| "subsampled dataset from {} to {} (ratio={})".format( |
| len(self.dataset), self.actual_size, size_ratio |
| ) |
| ) |
|
|
| def __getitem__(self, index): |
| return self.dataset[self.indices[index]] |
|
|
| def __len__(self): |
| return self.actual_size |
|
|
| def collater(self, samples): |
| return self.dataset.collater(samples) |
|
|
| @property |
| def sizes(self): |
| return self.dataset.sizes[self.indices] |
|
|
| @property |
| def name(self): |
| return self.dataset.name |
|
|
| def num_tokens(self, index): |
| return self.dataset.num_tokens(self.indices[index]) |
|
|
| def size(self, index): |
| return self.dataset.size(self.indices[index]) |
|
|
| def ordered_indices(self): |
| """Return an ordered list of indices. Batches will be constructed based |
| on this order.""" |
| if self.shuffle: |
| order = [np.random.permutation(len(self))] |
| else: |
| order = [np.arange(len(self))] |
| order.append(self.sizes) |
| return np.lexsort(order) |
|
|
| def prefetch(self, indices): |
| self.dataset.prefetch(self.indices[indices]) |
|
|
