Setting seed for Data Samplers (torch.utils.data.sampler)

Is it possible to assign a random seed for BatchSampler and SubsetRandomSampler? I am using a generator function (yields a generator object instead of returning it, as storing the data in memory turned out to be a lot slower). This function has to be called twice, while it is expected to return the same output.

I did my own random sampler which I can call set_seed() on every time I want to change the order of sampled data. If I’d want the same data twice, I’d call set_seed() after the second time I had read the data. Perhaps this could help you create your own sampler. data_source is my dataset.

class RandomSampler(Sampler):
  def __init__(self, data_source):
    self.data_source = data_source

  def set_seed(self):
    self.seed = random.randint(0, 2**32 - 1)

  def __iter__(self):
    n = len(self.data_source)
    indexes = list(range(n))
    random.Random(self.seed).shuffle(indexes)
    return iter(indexes)

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