def test_generator_getstate(test_case):
auto_gen = flow.Generator(device="auto")
state = auto_gen.get_state()
cpu_gen = flow.Generator(device="cpu")
state = cpu_gen.get_state()
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
cuda_gen = flow.Generator(device="cuda")
state = cuda_gen.get_state()
def test_different_devices(test_case):
auto_gen = flow.Generator(device="auto")
cpu_gen = flow.Generator(device="cpu")
test_case.assertTrue(auto_gen.initial_seed() == cpu_gen.initial_seed())
with test_case.assertRaises(RuntimeError) as context:
flow.Generator(device="invalid")
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
cuda_gen = flow.Generator(device="cuda")
test_case.assertTrue(auto_gen.initial_seed() == cuda_gen.initial_seed())
def test_generator_seed(test_case):
cpu_gen = flow.default_generator
auto_gen = flow.Generator(device="auto")
test_gens = [auto_gen, cpu_gen]
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
cuda_gen = flow.Generator(device="cuda")
cuda0_gen = flow.Generator(device="cuda:0")
test_gens += [cuda_gen, cuda0_gen]
for gen in test_gens:
seed = gen.seed()
test_case.assertTrue(seed == gen.initial_seed())
def test_generator_in_dropout(test_case):
tgt = flow.ones(2000000)
output = flow._C.dropout(tgt, 0.1, True, flow.Generator(), addend=None)
output.numpy()
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
output = flow._C.dropout(tgt.cuda(),
0.1,
True,
flow.Generator(),
addend=None)
output.numpy()
def test_generator_manual_seed(test_case):
cpu_gen = flow.default_generator
auto_gen = flow.Generator(device="auto")
test_gens = [cpu_gen, auto_gen]
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
cuda_gen = flow.Generator(device="cuda")
cuda0_gen = flow.Generator(device="cuda:0")
test_gens += [cuda_gen, cuda0_gen]
for seed in [1, 2]:
for gen in test_gens:
gen.manual_seed(seed)
test_case.assertTrue(gen.initial_seed() == seed)
def __iter__(self):
n = len(self.data_source)
if self.generator is None:
generator = flow.Generator("cpu")
generator.manual_seed(np.random.randint(0, np.iinfo(np.int64).max))
# TODO: use Tensor.random_
# generator.manual_seed(
# int(flow.empty((), dtype=flow.int64).random_().item())
# )
else:
generator = self.generator
if self.replacement:
for _ in range(self.num_samples // 32):
yield from flow._C.randint(
high=n, size=(32, ), dtype=flow.int64,
generator=generator).numpy().tolist()
yield from flow._C.randint(
high=n,
size=(self.num_samples % 32, ),
dtype=flow.int64,
generator=generator,
).numpy().tolist()
else:
yield from flow._C.randperm(n,
generator=generator).numpy().tolist()
def test_generator_setstate(test_case):
cpu_gen = flow.default_generator
flow.randn(100, 100, dtype=flow.float32, device="cpu", generator=cpu_gen)
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
cuda_gen = flow.Generator("cuda")
flow.randn(100, 100, dtype=flow.float32, device="cuda", generator=cuda_gen)
state = cpu_gen.get_state()
flow.randn(100, 100, dtype=flow.float32, device="cpu", generator=cpu_gen)
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
cuda_state = cuda_gen.get_state()
flow.randn(100, 100, dtype=flow.float32, device="cuda", generator=cuda_gen)
new_state = cpu_gen.get_state()
test_case.assertTrue(not np.allclose(new_state.numpy(), state.numpy()))
cpu_gen.set_state(state)
new_state = cpu_gen.get_state()
test_case.assertTrue(np.allclose(new_state.numpy(), state.numpy()))
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
new_cuda_state = cuda_gen.get_state()
test_case.assertTrue(
not np.allclose(new_cuda_state.numpy(), cuda_state.numpy())
)
cuda_gen.set_state(cuda_state)
new_cuda_state = cuda_gen.get_state()
test_case.assertTrue(
np.allclose(new_cuda_state.numpy(), cuda_state.numpy())
)
def __iter__(self) -> Iterator[T_co]:
if self.shuffle:
# deterministically shuffle based on epoch and seed
g = flow.Generator("cpu")
g.manual_seed(self.seed + self.epoch)
indices = flow._C.randperm(len(self.dataset), generator=g).tolist()
else:
indices = list(range(len(self.dataset)))
if not self.drop_last:
# add extra samples to make it evenly divisible
padding_size = self.total_size - len(indices)
if padding_size <= len(indices):
indices += indices[:padding_size]
else:
indices += (
indices *
math.ceil(padding_size / len(indices)))[:padding_size]
else:
# remove tail of data to make it evenly divisible.
indices = indices[:self.total_size]
assert len(indices) == self.total_size
# subsample
indices = indices[self.rank:self.total_size:self.num_replicas]
assert len(indices) == self.num_samples
return iter(indices)
def _rand_input(placement=None, sbp=None):
generator = flow.Generator()
generator.manual_seed(0)
return flow.randint(0,
10, (8, ),
generator=generator,
placement=placement,
sbp=sbp)
def _test_with_generator(test_case, shape, placement, sbp):
gen = flow.Generator()
gen.manual_seed(0)
y1 = flow.randn(*shape, placement=placement, sbp=sbp, generator=gen)
gen.manual_seed(0)
y2 = flow.randn(*shape, placement=placement, sbp=sbp, generator=gen)
test_case.assertTrue(
np.allclose(y1.numpy(), y2.numpy(), atol=1e-4, rtol=1e-4))
def _test_bernoulli_with_generator(test_case, shape):
generator = flow.Generator()
generator.manual_seed(0)
x = flow.tensor(np.random.rand(*shape),
dtype=flow.float32,
device=flow.device("cpu"))
y_1 = flow.bernoulli(x, generator=generator)
generator.manual_seed(0)
y_2 = flow.bernoulli(x, generator=generator)
test_case.assertTrue(np.allclose(y_1.numpy(), y_2.numpy()))
def _test_randperm_with_generator(test_case, N, device, dtype):
generator = flow.Generator()
generator.manual_seed(0)
y_1 = flow.randperm(N, device=device, dtype=dtype, generator=generator)
generator.manual_seed(0)
y_2 = flow.randperm(N, device=device, dtype=dtype, generator=generator)
test_case.assertTrue(np.allclose(y_1.numpy(), y_2.numpy()))
test_case.assertTrue(
y_1.device == flow.device(device) and y_2.device == flow.device(device)
)
test_case.assertTrue(y_1.dtype == dtype and y_2.dtype == dtype)
def _test_with_generator(test_case, device, shape):
gen = flow.Generator()
gen.manual_seed(0)
y1 = flow.randn(
*shape, dtype=flow.float32, device=flow.device(device), generator=gen
)
gen.manual_seed(0)
y2 = flow.randn(
*shape, dtype=flow.float32, device=flow.device(device), generator=gen
)
test_case.assertTrue(np.allclose(y1.numpy(), y2.numpy(), atol=1e-4, rtol=1e-4))
def fixed_cpu_seed_dropout_test(test_case):
gen1 = flow.Generator()
gen1.manual_seed(5)
dropped_array1 = np.array([
[0.000000, 0.000000, 1.333333],
[1.333333, 0.000000, 1.333333],
[1.333333, 1.333333, 1.333333],
]).astype(np.float32)
dropout1 = flow.nn.Dropout(p=0.25, generator=gen1)
x = flow.ones((3, 3), dtype=flow.float32)
out1 = dropout1(x)
test_case.assertTrue(
np.allclose(out1.numpy(), dropped_array1, atol=1e-4, rtol=1e-4))
gen2 = flow.Generator()
gen2.manual_seed(7)
dropout2 = flow.nn.Dropout(p=0.5, generator=gen2)
dropped_array2 = np.array([[0.0, 0.0, 2.0], [0.0, 0.0, 2.0],
[2.0, 0.0, 2.0]]).astype(np.float32)
out2 = dropout2(x)
test_case.assertTrue(
np.allclose(out2.numpy(), dropped_array2, atol=1e-4, rtol=1e-4))
def _test_with_generator(test_case, mean, std, shape, device, dtype):
dtype = type_name_to_flow_type[dtype]
gen = flow.Generator()
gen.manual_seed(0)
y1 = flow.normal(mean,
std,
*shape,
generator=gen,
dtype=dtype,
device=flow.device(device))
gen.manual_seed(0)
y2 = flow.normal(mean,
std,
*shape,
generator=gen,
dtype=dtype,
device=flow.device(device))
test_case.assertTrue(np.array_equal(y1.numpy(), y2.numpy()))
def test_different_devices(test_case):
auto_gen = flow.Generator(device="auto")
cpu_gen = flow.default_generator
with test_case.assertRaises(RuntimeError) as context:
flow.Generator(device="invalid")
flow.Generator(device="cpu:1000")
if not os.getenv("ONEFLOW_TEST_CPU_ONLY"):
with test_case.assertRaises(
oneflow._oneflow_internal.exception.Exception) as context:
flow.Generator(device="cuda:1000")
cuda_gen = flow.Generator(device="cuda")
cuda0_gen = flow.Generator(device="cuda:0")
def __init__(self, p: float = 0.5, inplace: bool = False, generator=None):
_DropoutNd.__init__(self, p, inplace)
self.p = p
if generator is None:
generator = flow.Generator()
self.generator = generator
def test_with_generator(test_case):
gen = flow.Generator("cuda")
x = flow.randn(2, 3, device="cuda", generator=gen)
test_case.assertEqual(x.device,
flow.device(f"cuda:{flow.env.get_rank()}"))
def test_generator_manual_seed(test_case):
generator = flow.Generator()
generator.manual_seed(1)
test_case.assertTrue(generator.initial_seed() == 1)
generator.manual_seed(2)
test_case.assertTrue(generator.initial_seed() == 2)
