set seed at multiprocessing iterator

Hiroshiba · May 15, 2021 · 893356c · 893356c
1 parent a395291
commit 893356c
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/pytorch_trainer/iterators/multiprocess_iterator.py b/pytorch_trainer/iterators/multiprocess_iterator.py
@@ -9,6 +9,7 @@
 
 import numpy
 import six
+import torch
 
 from pytorch_trainer.dataset import iterator
 from pytorch_trainer.iterators import _statemachine
@@ -554,6 +555,9 @@ def _fetch_setup(dataset, mem_size, mem_bulk):
     _fetch_mem_size = mem_size
     _fetch_mem_bulk = mem_bulk
 
+    numpy.random.seed(multiprocessing.current_process().pid)
+    torch.manual_seed(multiprocessing.current_process().pid)
+
 
 def _fetch_run(inputs):
     i, index = inputs