set numpy random state based on worker_id

https://tanelp.github.io/posts/a-bug-that-plagues-thousands-of-open-source-ml-projects/
zalandoresearch · Apr 11, 2021 · 499dae8 · 499dae8
1 parent 4e97483
commit 499dae8
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/pts/model/estimator.py b/pts/model/estimator.py
@@ -118,6 +118,7 @@ def train_model(
             num_workers=num_workers,
             prefetch_factor=prefetch_factor,
             pin_memory=True,
+            worker_init_fn=self._worker_init_fn,
             **kwargs,
         )
 
@@ -138,6 +139,7 @@ def train_model(
                 num_workers=num_workers,
                 prefetch_factor=prefetch_factor,
                 pin_memory=True,
+                worker_init_fn=self._worker_init_fn,
                 **kwargs,
             )
 
@@ -155,6 +157,10 @@ def train_model(
             ),
         )
 
+    @staticmethod
+    def _worker_init_fn(worker_id):
+        np.random.seed(np.random.get_state()[1][0] + worker_id)
+
     def train(
         self,
         training_data: Dataset,