def __iter__(self):
# deterministically shuffle based on epoch
g = torch.Generator()
g.manual_seed(self.epoch)
indices = list(torch.randperm(len(self.dataset), generator=g))
# add extra samples to make it evenly divisible
indices += indices[:(self.total_size - len(indices))]
assert len(indices) == self.total_size
# subsample
offset = self.num_samples * self.rank
indices = indices[offset:offset + self.num_samples]
assert len(indices) == self.num_samples
return iter(indices)
python类Generator()的实例源码
def __iter__(self):
# deterministically shuffle based on epoch
g = torch.Generator()
g.manual_seed(self.epoch)
indices = list(torch.randperm(len(self.dataset), generator=g))
# add extra samples to make it evenly divisible
indices += indices[:(self.total_size - len(indices))]
assert len(indices) == self.total_size
# subsample
offset = self.num_samples * self.rank
indices = indices[offset:offset + self.num_samples]
assert len(indices) == self.num_samples
return iter(indices)
def test_RNGStateAliasing(self):
# Fork the random number stream at this point
gen = torch.Generator()
gen.set_state(torch.get_rng_state())
self.assertEqual(gen.get_state(), torch.get_rng_state())
target_value = torch.rand(1000)
# Dramatically alter the internal state of the main generator
_ = torch.rand(100000)
forked_value = torch.rand(gen, 1000)
self.assertEqual(target_value, forked_value, 0, "RNG has not forked correctly.")
def test_RNGStateAliasing(self):
# Fork the random number stream at this point
gen = torch.Generator()
gen.set_state(torch.get_rng_state())
self.assertEqual(gen.get_state(), torch.get_rng_state())
target_value = torch.rand(1000)
# Dramatically alter the internal state of the main generator
_ = torch.rand(100000)
forked_value = torch.rand(gen, 1000)
self.assertEqual(target_value, forked_value, 0, "RNG has not forked correctly.")
def test_RNGStateAliasing(self):
# Fork the random number stream at this point
gen = torch.Generator()
gen.set_state(torch.get_rng_state())
self.assertEqual(gen.get_state(), torch.get_rng_state())
target_value = torch.rand(1000)
# Dramatically alter the internal state of the main generator
_ = torch.rand(100000)
forked_value = torch.rand(gen, 1000)
self.assertEqual(target_value, forked_value, 0, "RNG has not forked correctly.")
def test_RNGStateAliasing(self):
# Fork the random number stream at this point
gen = torch.Generator()
gen.set_state(torch.get_rng_state())
self.assertEqual(gen.get_state(), torch.get_rng_state())
target_value = torch.rand(1000)
# Dramatically alter the internal state of the main generator
_ = torch.rand(100000)
forked_value = torch.rand(1000, generator=gen)
self.assertEqual(target_value, forked_value, 0, "RNG has not forked correctly.")
def test_RNGStateAliasing(self):
# Fork the random number stream at this point
gen = torch.Generator()
gen.set_state(torch.get_rng_state())
self.assertEqual(gen.get_state(), torch.get_rng_state())
target_value = torch.rand(1000)
# Dramatically alter the internal state of the main generator
_ = torch.rand(100000)
forked_value = torch.rand(1000, generator=gen)
self.assertEqual(target_value, forked_value, 0, "RNG has not forked correctly.")