Seeder: Add option for seeding shared across workers

Author: tomsiadev

fwr13y/seeder/paddle.py

@@ -21,18 +21,31 @@
 
 class Seeder:
     def __init__(self, master_seed, ngpus, local_rank):
-
+        
+        self.master_seed_was_none = master_seed is None
+        if master_seed is None and local_rank == 0:
+            print('INFO: master_seed is None in seeder.init, random master_seed will be generated (different one for each worker).')
+        
         self.master_seed = (
             master_seed if master_seed is not None else generate_master_seed_randomly()
         )
         self.seed_gen = SeedGen(self.master_seed, ngpus, local_rank)
         self._ext_generators = []
+        self._ext_generators_shared = []
 
-    def register_generator(self, gen):
-        self._ext_generators.append(gen)
+    def register_generator(self, gen, shared=False):
+        if shared:
+            if self.master_seed_was_none:
+                raise Exception('master_seed was None during seeder.init, seeds shared among workers cannot be used.')
+            self._ext_generators_shared.append(gen)
+        else:
+            self._ext_generators.append(gen)
 
     def unregister_generator(self, gen):
-        self._ext_generators.remove(gen)
+        try:
+            self._ext_generators.remove(gen)
+        except ValueError:
+            self._ext_generators_shared.remove(gen)
 
     def reseed(self, task, epoch):
         seed = self.seed_gen(task, epoch)
@@ -43,6 +56,11 @@ def reseed(self, task, epoch):
         for generator in self._ext_generators:
             generator(seed)
 
+        if self._ext_generators_shared:
+            shared_seed = self.seed_gen(task, epoch, shared_seed=True)
+            for generator in self._ext_generators_shared:
+                generator(shared_seed)
+
 
 _seeder_run = None
 
@@ -62,11 +80,11 @@ def get_master_seed():
     return _seeder_run.master_seed
 
 
-def register_generator(gen):
+def register_generator(gen, shared=False):
     global _seeder_run
-    _seeder_run._ext_generators.append(gen)
+    _seeder_run.register_generator(gen, shared)
 
 
 def unregister_generator(gen):
     global _seeder_run
-    _seeder_run._ext_generators.remove(gen)
+    _seeder_run.unregister_generator(gen)

fwr13y/seeder/pyt.py

@@ -21,18 +21,31 @@
 
 class Seeder:
     def __init__(self, master_seed, ngpus, local_rank):
-
+        
+        self.master_seed_was_none = master_seed is None
+        if master_seed is None and local_rank == 0:
+            print('INFO: master_seed is None in seeder.init, random master_seed will be generated (different one for each worker).')
+        
         self.master_seed = (
             master_seed if master_seed is not None else generate_master_seed_randomly()
         )
         self.seed_gen = SeedGen(self.master_seed, ngpus, local_rank)
         self._ext_generators = []
+        self._ext_generators_shared = []
 
-    def register_generator(self, gen):
-        self._ext_generators.append(gen)
+    def register_generator(self, gen, shared=False):
+        if shared:
+            if self.master_seed_was_none:
+                raise Exception('master_seed was None during seeder.init, seeds shared among workers cannot be used.')
+            self._ext_generators_shared.append(gen)
+        else:
+            self._ext_generators.append(gen)
 
     def unregister_generator(self, gen):
-        self._ext_generators.remove(gen)
+        try:
+            self._ext_generators.remove(gen)
+        except ValueError:
+            self._ext_generators_shared.remove(gen)
 
     def reseed(self, task, epoch):
         seed = self.seed_gen(task, epoch)
@@ -43,6 +56,11 @@ def reseed(self, task, epoch):
         for generator in self._ext_generators:
             generator(seed)
 
+        if self._ext_generators_shared:
+            shared_seed = self.seed_gen(task, epoch, shared_seed=True)
+            for generator in self._ext_generators_shared:
+                generator(shared_seed)
+
 
 _seeder_run = None
 
@@ -62,11 +80,11 @@ def get_master_seed():
     return _seeder_run.master_seed
 
 
-def register_generator(gen):
+def register_generator(gen, shared=False):
     global _seeder_run
-    _seeder_run._ext_generators.append(gen)
+    _seeder_run.register_generator(gen, shared)
 
 
 def unregister_generator(gen):
     global _seeder_run
-    _seeder_run._ext_generators.remove(gen)
+    _seeder_run.unregister_generator(gen)

fwr13y/seeder/seed_gen.py

@@ -29,10 +29,19 @@ def __init__(self, master_seed, ngpus, local_rank):
         self._used_seeds = set()
         self._rng = random.Random(0)
 
-    def __call__(self, task, epoch):
-        seed = (
-            self.master_seed + (epoch * self.ngpus + self.local_rank)
-        ) * self.ntasks + task
+    def __call__(self, task, epoch, shared_seed=False):
+        if shared_seed:
+            # Use the same seed for every rank
+            # Constant at the beginning so that the values do not repeat with not shared seeds
+            seed = 2 * (
+                self.master_seed + epoch
+            ) * self.ntasks + task
+        else:
+            # Use different seed for every rank
+            seed = (
+                self.master_seed + (epoch * self.ngpus + self.local_rank)
+            ) * self.ntasks + task
+
         if seed in self._used_seeds:
             print(
                 "Warning!!! seed has been generated more than once!!!", file=sys.stderr

fwr13y/seeder/tf.py

@@ -21,18 +21,31 @@
 
 class Seeder:
     def __init__(self, master_seed, ngpus, local_rank):
-
+        
+        self.master_seed_was_none = master_seed is None
+        if master_seed is None and local_rank == 0:
+            print('INFO: master_seed is None in seeder.init, random master_seed will be generated (different one for each worker).')
+        
         self.master_seed = (
             master_seed if master_seed is not None else generate_master_seed_randomly()
         )
         self.seed_gen = SeedGen(self.master_seed, ngpus, local_rank)
         self._ext_generators = []
+        self._ext_generators_shared = []
 
-    def register_generator(self, gen):
-        self._ext_generators.append(gen)
+    def register_generator(self, gen, shared=False):
+        if shared:
+            if self.master_seed_was_none:
+                raise Exception('master_seed was None during seeder.init, seeds shared among workers cannot be used.')
+            self._ext_generators_shared.append(gen)
+        else:
+            self._ext_generators.append(gen)
 
     def unregister_generator(self, gen):
-        self._ext_generators.remove(gen)
+        try:
+            self._ext_generators.remove(gen)
+        except ValueError:
+            self._ext_generators_shared.remove(gen)
 
     def reseed(self, task, epoch):
         seed = self.seed_gen(task, epoch)
@@ -43,6 +56,11 @@ def reseed(self, task, epoch):
         for generator in self._ext_generators:
             generator(seed)
 
+        if self._ext_generators_shared:
+            shared_seed = self.seed_gen(task, epoch, shared_seed=True)
+            for generator in self._ext_generators_shared:
+                generator(shared_seed)
+
 
 _seeder_run = None
 
@@ -62,15 +80,14 @@ def get_master_seed():
     return _seeder_run.master_seed
 
 
-def register_generator(gen):
+def register_generator(gen, shared=False):
     global _seeder_run
-    _seeder_run._ext_generators.append(gen)
+    _seeder_run.register_generator(gen, shared)
 
 
 def unregister_generator(gen):
     global _seeder_run
-    _seeder_run._ext_generators.remove(gen)
-
+    _seeder_run.unregister_generator(gen)
 
 class SeederCB(tf.keras.callbacks.Callback):
     def on_epoch_begin(self, epoch, logs=None):