Dev refine generator (#5426)

* add bernoulli module

* fix doc test

* add bernoulli functor

* make changes according to review

* refine (#5415)

* refine rng related tests; make generator.seed() as api

Co-authored-by: Ldpe2G <liangdepeng@gmail.com>
Co-authored-by: oneflow-ci-bot <69100618+oneflow-ci-bot@users.noreply.github.com>

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) {

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_; }

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_;

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__":

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",

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"]