Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,000 Bytes
05fb4ab |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 |
import torch
from torch.utils.data import Sampler
class RandomSampler(Sampler):
"""Randomly samples items (num_samples) at each epoch. """
def __init__(self, data_source, num_samples=None):
super().__init__(data_source)
self.data_source = data_source
self._num_samples = num_samples
if not isinstance(self.num_samples, int) or self.num_samples <= 0:
raise ValueError(
"num_samples should be a positive integer "
"value, but got num_samples={}".format(self.num_samples)
)
@property
def num_samples(self):
# dataset load_size might change at runtime
if self._num_samples is None:
return len(self.data_source)
return self._num_samples
def __iter__(self):
n = len(self.data_source)
new_list = torch.randperm(n, dtype=torch.int64)[: self.num_samples].tolist()
return iter(new_list)
def __len__(self):
return self.num_samples
class FirstItemsSampler(Sampler):
"""Samples the first 'num_samples' iterms at each epoch. Useful for degubbing. """
def __init__(self, data_source, num_samples=None):
super().__init__(data_source)
self.data_source = data_source
self._num_samples = num_samples
if not isinstance(self.num_samples, int) or self.num_samples <= 0:
raise ValueError(
"num_samples should be a positive integer "
"value, but got num_samples={}".format(self.num_samples)
)
@property
def num_samples(self):
# dataset load_size might change at runtime
if self._num_samples is None:
return len(self.data_source)
return self._num_samples
def __iter__(self):
n = len(self.data_source)
new_list = torch.arange(start=0, end=n, dtype=torch.int64)[: self.num_samples].tolist()
return iter(new_list)
def __len__(self):
return self.num_samples |