diff --git a/oneflow/api/python/framework/random_generator.cpp b/oneflow/api/python/framework/random_generator.cpp
index bded13ecf4a5f159b32ff280c5c27fd545480a5f..31d12bb9ad60e62c332612f20af1b743bc853f74 100644
--- a/oneflow/api/python/framework/random_generator.cpp
+++ b/oneflow/api/python/framework/random_generator.cpp
@@ -26,7 +26,8 @@ namespace oneflow {
ONEFLOW_API_PYBIND11_MODULE("", m) {
py::class_<one::Generator, std::shared_ptr<one::Generator>>(m, "Generator")
.def("manual_seed", &one::Generator::set_current_seed)
- .def("initial_seed", &one::Generator::current_seed);
+ .def("initial_seed", &one::Generator::current_seed)
+ .def("seed", &one::Generator::seed);
m.def("manual_seed", [](uint64_t seed) { return one::ManualSeed(seed).GetOrThrow(); });
m.def("create_generator", [](const std::string& device_tag) {
diff --git a/oneflow/core/framework/random_generator.h b/oneflow/core/framework/random_generator.h
index c7074c27edc647f658cc440667a429f91835334c..3e6ae9f9c48363bf502addcb9568e50f354cfa77 100644
--- a/oneflow/core/framework/random_generator.h
+++ b/oneflow/core/framework/random_generator.h
@@ -36,7 +36,7 @@ class Generator final {
uint64_t current_seed() const;
- // Reset current seed by default seed, and returns it.
+ // Reset current generator by a non-deterministic random seed, and returns it.
uint64_t seed();
const std::shared_ptr<GeneratorImpl>& impl() const { return impl_; }
diff --git a/oneflow/core/framework/random_generator_impl.h b/oneflow/core/framework/random_generator_impl.h
index e53abc59a9fa651693ed643b96dd6484e14013f6..666e57e119ea798e43d3d2439e4c3279f994ea71 100644
--- a/oneflow/core/framework/random_generator_impl.h
+++ b/oneflow/core/framework/random_generator_impl.h
@@ -94,7 +94,6 @@ class CPUGeneratorImpl : public DeviceGeneratorImpl {
}
std::mt19937& engine() { return engine_; }
- void set_engine(std::mt19937 engine) { engine_ = engine; }
public:
std::mt19937 engine_;
diff --git a/oneflow/python/test/generator/test_generator.py b/oneflow/python/test/generator/test_generator.py
index b61bced5a1cb0f42d272aab92af8a83df18abf95..94294360edea0c8a47b23f0a5153947c7c556343 100644
--- a/oneflow/python/test/generator/test_generator.py
+++ b/oneflow/python/test/generator/test_generator.py
@@ -44,19 +44,16 @@ class TestDefaultGenerator(flow.unittest.TestCase):
cuda_gen = flow.default_generator(device="cuda")
cuda0_gen = flow.default_generator(device="cuda:0")
cpu_gen = flow.default_generator(device="cpu")
- test_case.assertTrue(auto_gen.initial_seed() == global_seed)
- test_case.assertTrue(cuda_gen.initial_seed() == global_seed)
- test_case.assertTrue(cuda0_gen.initial_seed() == global_seed)
- test_case.assertTrue(cpu_gen.initial_seed() == global_seed)
+ for gen in [auto_gen, cuda_gen, cuda0_gen, cpu_gen]:
+ test_case.assertTrue(gen.initial_seed() == global_seed)
def test_different_devices(test_case):
auto_gen = flow.default_generator(device="auto")
cuda_gen = flow.default_generator(device="cuda")
cuda0_gen = flow.default_generator(device="cuda:0")
cpu_gen = flow.default_generator(device="cpu")
- test_case.assertTrue(auto_gen.initial_seed() == cuda_gen.initial_seed())
- test_case.assertTrue(auto_gen.initial_seed() == cuda0_gen.initial_seed())
- test_case.assertTrue(auto_gen.initial_seed() == cpu_gen.initial_seed())
+ for gen in [cuda_gen, cuda0_gen, cpu_gen]:
+ test_case.assertTrue(auto_gen.initial_seed() == gen.initial_seed())
with test_case.assertRaises(Exception) as context:
flow.default_generator(device="invalid")
@@ -75,15 +72,19 @@ class TestDefaultGenerator(flow.unittest.TestCase):
cuda_gen = flow.default_generator(device="cuda")
cpu_gen = flow.default_generator(device="cpu")
- auto_gen.manual_seed(1)
- test_case.assertTrue(auto_gen.initial_seed() == 1)
- test_case.assertTrue(cpu_gen.initial_seed() == 1)
- test_case.assertTrue(cuda_gen.initial_seed() == 1)
+ for seed in [1, 2]:
+ auto_gen.manual_seed(seed)
+ for gen in [auto_gen, cuda_gen, cpu_gen]:
+ test_case.assertTrue(gen.initial_seed() == seed)
- auto_gen.manual_seed(2)
- test_case.assertTrue(auto_gen.initial_seed() == 2)
- test_case.assertTrue(cuda_gen.initial_seed() == 2)
- test_case.assertTrue(cpu_gen.initial_seed() == 2)
+ def test_generator_seed(test_case):
+ auto_gen = flow.default_generator(device="auto")
+ cuda_gen = flow.default_generator(device="cuda")
+ cpu_gen = flow.default_generator(device="cpu")
+
+ for gen in [auto_gen, cuda_gen, cpu_gen]:
+ seed = gen.seed()
+ test_case.assertTrue(seed == gen.initial_seed())
if __name__ == "__main__":
diff --git a/oneflow/python/test/modules/test_bernoulli.py b/oneflow/python/test/modules/test_bernoulli.py
index 8b80cfa50c3b8e910d4d65a4dacf466aecbbdb16..30eed8270d2e1dd75ee34ad6455633894a729580 100644
--- a/oneflow/python/test/modules/test_bernoulli.py
+++ b/oneflow/python/test/modules/test_bernoulli.py
@@ -29,6 +29,17 @@ def _test_bernoulli(test_case, shape):
test_case.assertTrue(np.allclose(y.numpy(), x.numpy()))
+def _test_bernoulli_with_generator(test_case, shape):
+ generator = flow.Generator()
+ generator.manual_seed(0)
+ x = flow.Tensor(np.random.rand(*shape), device=flow.device("cpu"))
+ y_1 = flow.bernoulli(x, generator=generator)
+ y_1.numpy() # sync
+ generator.manual_seed(0)
+ y_2 = flow.bernoulli(x, generator=generator)
+ test_case.assertTrue(np.allclose(y_1.numpy(), y_2.numpy()))
+
+
@unittest.skipIf(
not flow.unittest.env.eager_execution_enabled(),
".numpy() doesn't work in lazy mode",
diff --git a/oneflow/python/test/modules/test_dropout.py b/oneflow/python/test/modules/test_dropout.py
index 7f279cd2b69c7521e03ae3b4f56ad97439dcc701..2029ffb6a0a081cd5732c41e658304beb4f47065 100644
--- a/oneflow/python/test/modules/test_dropout.py
+++ b/oneflow/python/test/modules/test_dropout.py
@@ -77,6 +77,18 @@ def _test_dropout_eval(test_case, shape, device):
test_case.assertTrue(np.allclose(y.numpy(), input_arr))
+def _test_dropout_with_generator(test_case, shape, device):
+ generator = flow.Generator()
+ generator.manual_seed(0)
+ m = flow.nn.Dropout(p=0.5, generator=generator)
+ x = flow.Tensor(np.random.randn(*shape), device=flow.device(device))
+ y_1 = m(x)
+ y_1.numpy() # sync
+ generator.manual_seed(0)
+ y_2 = m(x)
+ test_case.assertTrue(np.allclose(y_1.numpy(), y_2.numpy()))
+
+
@unittest.skipIf(
not flow.unittest.env.eager_execution_enabled(),
".numpy() doesn't work in lazy mode",
@@ -90,6 +102,7 @@ class TestDropout(flow.unittest.TestCase):
_test_dropout_backward_p0,
_test_dropout_backward_p1,
_test_dropout_eval,
+ _test_dropout_with_generator,
]
arg_dict["shape"] = [(2, 3), (2, 3, 4), (2, 3, 4, 5)]
arg_dict["device"] = ["cpu", "cuda"]