Repository: pytorch/pytorch
Pull request: Add default PyTorch seeding and worker_init_fn to DataLoader
Fixes #3880 .

Problem statement:
Reinitialize the random generator in worker processes of `DataLoader`
`torch.utils.data.DataLoader` can use multiprocessing to load and preprocess the data. It is commonly used to overlap the GPU computation and data loading. It has a minor flaw, though. On Unix, the processes are by default created with `fork`, so the global random state is copied. If each worker does something pseudo-randomly (common in data augmentation), their actions are all the same. For example, random crop is commonly used in image classification task, and in this case, all `DataLoader` workers will crop image in the same way, reducing the effectiveness of data augmentation.

I propose a new option `worker_init_fn` be added to `DataLoader` constructor, and the function is called with the process ID right in the beginning of each worker loop. By default, the function ignores the process ID and reseeds both pytorch and numpy default random generator. The users can override it to do other things that each worker might need differently.

Hints from issue discussion:
this seems like a reasonable solution,. though needs some thought on the interface. We cant rely on `set_start_method` all the time right. wdyt @colesbury ?
Yes, this seems reasonable. @apaszke suggested using the worker ID (i.e. 0, 1, ... N-1) instead of PID to preserve determinism when the user sets the torch.manual_seed.

I think you probably want to combine the worker ID with some function of the current RNG state so that changing the seed in the main process changes affects the seeds in the worker processes.
I tried to reproduce this issue and there is an interesting behaviour (at least for me) that I observe. If random data augmentation is based on `np.random` then forking the dataset with random transformations produces batches with exactly the same augmentations. However, if random data augmentation is based on `random` from python stdlib there is no problem with data augmentations between batches. [Here](https://gist.github.com/vfdev-5/aa61ea2bac2494ca852963fc941fcd7e) is a gist.
@vfdev-5 that is because np.random doesn't work nicely with multiprocessing without manual seeding.
looking into this.

Commit message: Add default PyTorch seeding and worker_init_fn to DataLoader

Patch to review:
diff --git a/test/test_dataloader.py b/test/test_dataloader.py
index 8a64af507b03..d3197d6a8ca9 100644
--- a/test/test_dataloader.py
+++ b/test/test_dataloader.py
@@ -112,6 +112,18 @@ def __len__(self):
         return self.size
 
 
+class SeedDataset(Dataset):
+
+    def __init__(self, size):
+        self.size = size
+
+    def __getitem__(self, idx):
+        return torch.initial_seed()
+
+    def __len__(self):
+        return self.size
+
+
 class TestDataLoader(TestCase):
 
     def setUp(self):
@@ -220,6 +232,23 @@ def _test_timeout():
         finally:
             p.terminate()
 
+    def test_worker_seeding(self):
+        # worker has seed (main_proc_seed + worker_id + 1)
+        dataset = SeedDataset(4)
+        dataloader = DataLoader(dataset, batch_size=1, num_workers=1)
+        seed = torch.initial_seed()
+        batch = next(iter(dataloader))
+        self.assertEqual(seed + 1, batch[0])
+
+    def test_worker_init_fn(self):
+        # test custom init function
+        dataset = SeedDataset(4)
+        dataloader = DataLoader(dataset, batch_size=2, num_workers=2,
+                                worker_init_fn=lambda i: torch.manual_seed(12345))
+        for batch in dataloader:
+            self.assertEqual(12345, batch[0])
+            self.assertEqual(12345, batch[1])
+
     def test_shuffle(self):
         self._test_shuffle(DataLoader(self.dataset, shuffle=True))
 
diff --git a/torch/utils/data/dataloader.py b/torch/utils/data/dataloader.py
index 8aaf72a383db..358645a5ed4c 100644
--- a/torch/utils/data/dataloader.py
+++ b/torch/utils/data/dataloader.py
@@ -4,6 +4,7 @@
     _remove_worker_pids, _error_if_any_worker_fails
 from .sampler import SequentialSampler, RandomSampler, BatchSampler
 import signal
+import functools
 import collections
 import re
 import sys
@@ -30,7 +31,7 @@ def __init__(self, exc_info):
         self.exc_msg = "".join(traceback.format_exception(*exc_info))
 
 
-def _worker_loop(dataset, index_queue, data_queue, collate_fn):
+def _worker_loop(dataset, index_queue, data_queue, collate_fn, init_fn, id):
     global _use_shared_memory
     _use_shared_memory = True
 
@@ -41,6 +42,11 @@ def _worker_loop(dataset, index_queue, data_queue, collate_fn):
     _set_worker_signal_handlers()
 
     torch.set_num_threads(1)
+    torch.manual_seed(torch.initial_seed() + id + 1)
+
+    if init_fn is not None:
+        init_fn(id)
+
     while True:
         r = index_queue.get()
         if r is None:
@@ -183,6 +189,7 @@ def __init__(self, loader):
         self.sample_iter = iter(self.batch_sampler)
 
         if self.num_workers > 0:
+            self.worker_init_fn = loader.worker_init_fn
             self.index_queue = multiprocessing.SimpleQueue()
             self.worker_result_queue = multiprocessing.SimpleQueue()
             self.batches_outstanding = 0
@@ -195,8 +202,9 @@ def __init__(self, loader):
             self.workers = [
                 multiprocessing.Process(
                     target=_worker_loop,
-                    args=(self.dataset, self.index_queue, self.worker_result_queue, self.collate_fn))
-                for _ in range(self.num_workers)]
+                    args=(self.dataset, self.index_queue, self.worker_result_queue, self.collate_fn,
+                          self.worker_init_fn, i))
+                for i in range(self.num_workers)]
 
             for w in self.workers:
                 w.daemon = True  # ensure that the worker exits on process exit
@@ -326,7 +334,7 @@ class DataLoader(object):
             indices at a time. Mutually exclusive with batch_size, shuffle,
             sampler, and drop_last.
         num_workers (int, optional): how many subprocesses to use for data
-            loading. 0 means that the data will be loaded in the main process
+            loading. 0 means that the data will be loaded in the main process.
             (default: 0)
         collate_fn (callable, optional): merges a list of samples to form a mini-batch.
         pin_memory (bool, optional): If ``True``, the data loader will copy tensors
@@ -337,11 +345,19 @@ class DataLoader(object):
             will be smaller. (default: False)
         timeout (numeric, optional): if positive, the timeout value for collecting a batch
             from workers. Should always be non-negative. (default: 0)
+        worker_init_fn (callable, optional): If not None, this will be called on each
+            worker subprocess with the worker id after seeding and before data loading.
+            (default: None)
+
+    .. note: By default, the workers will have their PyTorch seed set before data loading
+             as `(main_process_seed + worker_id + 1)`, i.e.,
+             `torch.manual_seed(torch.initial_seed() + id + 1)`. If setting other
+             seeds (e.g. NumPy) is desired, it can be done through :attr:`worker_init_fn`.
     """
 
     def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sampler=None,
                  num_workers=0, collate_fn=default_collate, pin_memory=False, drop_last=False,
-                 timeout=0):
+                 timeout=0, worker_init_fn=None):
         self.dataset = dataset
         self.batch_size = batch_size
         self.num_workers = num_workers
@@ -349,6 +365,7 @@ def __init__(self, dataset, batch_size=1, shuffle=False, sampler=None, batch_sam
         self.pin_memory = pin_memory
         self.drop_last = drop_last
         self.timeout = timeout
+        self.worker_init_fn = worker_init_fn
 
         if timeout < 0:
             raise ValueError('timeout option should be non-negative')


Code review comments:
diff_hunk: @@ -30,7 +31,7 @@ def __init__(self, exc_info):
         self.exc_msg = "".join(traceback.format_exception(*exc_info))
 
 
-def _worker_loop(dataset, index_queue, data_queue, collate_fn):
+def _worker_loop(dataset, index_queue, data_queue, collate_fn, init_fn, id):
path: torch/utils/data/dataloader.py
text: @user1:
Please don't use `id`, because you're shading the `id` builtin, which is quite useful. `worker_id` would be better

diff_hunk: @@ -195,8 +202,9 @@ def __init__(self, loader):
             self.workers = [
                 multiprocessing.Process(
                     target=_worker_loop,
-                    args=(self.dataset, self.index_queue, self.worker_result_queue, self.collate_fn))
-                for _ in range(self.num_workers)]
+                    args=(self.dataset, self.index_queue, self.worker_result_queue, self.collate_fn,
+                          self.worker_init_fn, i))
path: torch/utils/data/dataloader.py
text: @user1:
Can you check if this works with `spawn` start method? I'm not sure if you can serialize arbitrary lambdas like this.