Auto-set DataLoader.worker_init_fn with seed_everything

pytorch_lightning/trainer/data_loading.py

@@ -13,8 +13,11 @@
 # limitations under the License.
 import inspect
 import multiprocessing
+import os
 from abc import ABC
 from copy import deepcopy
+
+from functools import partial
 from typing import Iterable, List, Optional, Tuple, Union
 
 from torch.utils.data import BatchSampler, DataLoader, RandomSampler, SequentialSampler
@@ -30,6 +33,7 @@
 from pytorch_lightning.utilities.debugging import InternalDebugger
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
 from pytorch_lightning.utilities.model_helpers import is_overridden
+from pytorch_lightning.utilities.seed import pl_worker_init_function
 
 
 class TrainerDataLoadingMixin(ABC):
@@ -100,6 +104,10 @@ def _worker_check(self, dataloader: DataLoader, name: str) -> None:
                 f' in the `DataLoader` init to improve performance.'
             )
 
+    def auto_add_worker_init_fn(self, dataloader: DataLoader) -> None:
+        if dataloader.worker_init_fn is None and int(os.environ.get("PL_SEED_WORKERS", "0")):
+            dataloader.worker_init_fn = partial(pl_worker_init_function, rank=self.global_rank)
+
     def auto_add_sampler(self, dataloader: DataLoader, shuffle: bool) -> DataLoader:
 
         # don't do anything if it's not a dataloader
@@ -231,6 +239,9 @@ def reset_train_dataloader(self, model: LightningModule) -> None:
         # check the workers recursively
         apply_to_collection(self.train_dataloader, DataLoader, self._worker_check, 'train dataloader')
 
+        # add worker_init_fn for correct seeding in worker processes
+        apply_to_collection(self.train_dataloader, DataLoader, self.auto_add_worker_init_fn)
+
         # wrap the sequence of train loaders to a CombinedLoader object for computing the num_training_batches
         self.train_dataloader = CombinedLoader(self.train_dataloader, self._multiple_trainloader_mode)
 
@@ -329,6 +340,9 @@ def _reset_eval_dataloader(
         # add samplers
         dataloaders = [self.auto_add_sampler(dl, shuffle=False) for dl in dataloaders if dl is not None]
 
+        # add worker_init_fn for correct seeding in worker processes
+        apply_to_collection(dataloaders, dtype=DataLoader, function=self.auto_add_worker_init_fn)
+
         loader_num_batches = []
 
         # determine number of batches

pytorch_lightning/utilities/seed.py

@@ -20,23 +20,30 @@
 
 import numpy as np
 import torch
+from pytorch_lightning.utilities.distributed import rank_zero_only
 
 from pytorch_lightning.utilities import rank_zero_warn
 
 log = logging.getLogger(__name__)
 
 
-def seed_everything(seed: Optional[int] = None) -> int:
+def seed_everything(seed: Optional[int] = None, workers: bool = True) -> int:
     """
     Function that sets seed for pseudo-random number generators in:
     pytorch, numpy, python.random
-    In addition, sets the env variable `PL_GLOBAL_SEED` which will be passed to
-    spawned subprocesses (e.g. ddp_spawn backend).
+    In addition, sets the following environment variables:
+
+    - `PL_GLOBAL_SEED`: will be passed to spawned subprocesses (e.g. ddp_spawn backend).
+    - `PL_SEED_WORKERS`: (optional) is set to 1 if ```workers=True``.
 
     Args:
         seed: the integer value seed for global random state in Lightning.
             If `None`, will read seed from `PL_GLOBAL_SEED` env variable
             or select it randomly.
+        workers: if set to ``True``, will properly configure all dataloaders passed to the
+            Trainer with a ``worker_init_fn``. If the user already provides such a function
+            for their dataloaders, setting this argument will have no influence. See also:
+            :func:`~pytorch_lightning.utilities.seed.pl_worker_init_function`.
     """
     max_seed_value = np.iinfo(np.uint32).max
     min_seed_value = np.iinfo(np.uint32).min
@@ -61,8 +68,35 @@ def seed_everything(seed: Optional[int] = None) -> int:
     np.random.seed(seed)
     torch.manual_seed(seed)
     torch.cuda.manual_seed_all(seed)
+
+    os.environ["PL_SEED_WORKERS"] = f"{int(workers)}"
+
     return seed
 
 
 def _select_seed_randomly(min_seed_value: int = 0, max_seed_value: int = 255) -> int:
     return random.randint(min_seed_value, max_seed_value)
+
+
+def pl_worker_init_function(worker_id: int, rank: Optional = None) -> None:
+    """
+    The worker_init_fn that Lightning automatically adds to your dataloader if you previously set
+    set the seed with :func:`~pytorch_lightning.utilities.seed.seed_everything`.
+    See also the PyTorch documentation on
+    `randomness in DataLoaders <https://pytorch.org/docs/stable/notes/randomness.html#dataloader>`_.
+    """
+    # implementation notes: https://github.com/pytorch/pytorch/issues/5059#issuecomment-817392562
+    global_rank = rank if rank is not None else rank_zero_only.rank
+    process_seed = torch.initial_seed()
+    # back out the base seed so we can use all the bits
+    base_seed = process_seed - worker_id
+    ss = np.random.SeedSequence([base_seed, worker_id, global_rank])
+    # use 128 bits (4 x 32-bit words)
+    np.random.seed(ss.generate_state(4))
+    # Spawn distinct SeedSequences for the PyTorch PRNG and the stdlib random module
+    torch_ss, stdlib_ss = ss.spawn(2)
+    # PyTorch takes a 64-bit seed
+    torch.manual_seed(torch_ss.generate_state(1, dtype=np.uint64)[0])
+    # use 128 bits expressed as an integer
+    stdlib_seed = (stdlib_ss.generate_state(2, dtype=np.uint64).astype(object) * [1 << 64, 1]).sum()
+    random.seed(stdlib_seed)

tests/trainer/test_data_loading.py

@@ -11,7 +11,6 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-
 import pytest
 from torch.utils.data import DataLoader
 from torch.utils.data.sampler import BatchSampler, SequentialSampler
@@ -72,7 +71,7 @@ def test_dataloader(self):
         return [self.create_dataset()] * self._numbers_test_dataloaders
 
 
-def check_replace_distrubuted_sampler(tmpdir, save_preds_on_dl_idx, accelerator, gpus, num_dl_idx, mode):
+def check_replace_distributed_sampler(tmpdir, save_preds_on_dl_idx, accelerator, gpus, num_dl_idx, mode):
     num_processes = 2
     limit_test_batches = 2
     trainer_args = {
@@ -100,8 +99,8 @@ def check_replace_distrubuted_sampler(tmpdir, save_preds_on_dl_idx, accelerator,
 
 @RunIf(min_gpus=2, special=True)
 @pytest.mark.parametrize("mode", [1, 2])
-def test_replace_distrubuted_sampler_custom_dataloader_custom_batch_sampler(tmpdir, mode):
-    check_replace_distrubuted_sampler(tmpdir, True, "ddp", 2, 2, mode)
+def test_replace_distributed_sampler_custom_dataloader_custom_batch_sampler(tmpdir, mode):
+    check_replace_distributed_sampler(tmpdir, True, "ddp", 2, 2, mode)
 
 
 @pytest.mark.parametrize("num_workers", [0, 1])

tests/trainer/test_dataloaders.py

@@ -12,21 +12,25 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import os
+from functools import partial
 from unittest import mock
-from unittest.mock import patch
+from unittest.mock import patch, Mock
 
+import numpy
 import pytest
 import torch
 from torch.utils.data.dataloader import DataLoader
-from torch.utils.data.dataset import IterableDataset, Subset
+from torch.utils.data.dataset import IterableDataset, Subset, Dataset
 from torch.utils.data.distributed import DistributedSampler
 from torch.utils.data.sampler import SequentialSampler
 
 import tests.helpers.pipelines as tpipes
-from pytorch_lightning import Callback, Trainer
+from pytorch_lightning import Callback, Trainer, seed_everything
 from pytorch_lightning.trainer.states import TrainerState
+from pytorch_lightning.utilities import rank_zero_only
 from pytorch_lightning.utilities.data import has_iterable_dataset, has_len
 from pytorch_lightning.utilities.exceptions import MisconfigurationException
+from pytorch_lightning.utilities.seed import pl_worker_init_function
 from tests.base import EvalModelTemplate
 from tests.helpers.boring_model import BoringModel, RandomDataset
 from tests.helpers.runif import RunIf
@@ -634,6 +638,88 @@ def test_warning_with_few_workers_multi_loader(_, tmpdir, ckpt_path, stage):
             trainer.fit(model, train_dataloader=train_multi_dl, val_dataloaders=val_multi_dl)
 
 
+def _user_worker_init_fn(_):
+    pass
+
+
+def test_auto_add_worker_init_fn():
+    """ Test Trainer adds a default worker_init_fn to the dataloader when seed_everything() is used. """
+    dataset = Mock()
+    dataloader = DataLoader(dataset)
+    trainer = Trainer()
+
+    # without pl.seed_everything()
+    trainer.auto_add_worker_init_fn(dataloader)
+    assert dataloader.worker_init_fn is None
+
+    # with forcefully avoiding it
+    seed_everything(0, workers=False)
+    trainer.auto_add_worker_init_fn(dataloader)
+    assert dataloader.worker_init_fn is None
+
+    # when user already has a worker_init_fn
+    user_function = _user_worker_init_fn
+    dataloader.worker_init_fn = user_function
+    trainer.auto_add_worker_init_fn(dataloader)
+    assert dataloader.worker_init_fn is user_function
+    dataloader.worker_init_fn = None
+
+    # main use case
+    seed_everything(0, workers=True)
+    trainer.auto_add_worker_init_fn(dataloader)
+    assert dataloader.worker_init_fn is not None
+
+
+class NumpyRandomDataset(Dataset):
+    size = 16
+
+    def __getitem__(self, index):
+        return numpy.random.randint(0, 100, 3)
+
+    def __len__(self):
+        return self.size
+
+
+class MultiProcessModel(BoringModel):
+
+    def __init__(self):
+        super().__init__()
+        self.batches_seen = []
+
+    def training_step(self, batch, batch_idx):
+        self.batches_seen.append(batch)
+
+    def training_epoch_end(self, outputs):
+        world_size = 2
+        num_samples = NumpyRandomDataset.size
+        all_batches = torch.cat(self.batches_seen)
+        all_batches = self.all_gather(all_batches)
+        assert all_batches.shape[0] == world_size
+        all_batches = all_batches.view(-1, 3)
+        assert len(torch.unique(all_batches, dim=0)) == num_samples
+
+
+@RunIf(min_gpus=2)
+def test_auto_add_worker_init_fn_distributed(tmpdir, monkeypatch):
+    """ Test that the lightning worker_init_fn takes care of dataloaders in multi-gpu/multi-node training. """
+    dataset = NumpyRandomDataset()
+    num_workers = 2
+    batch_size = 2
+
+    dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=num_workers)
+    seed_everything(0, workers=True)
+    trainer = Trainer(
+        default_root_dir=tmpdir,
+        max_epochs=1,
+        gpus=2,
+        accelerator="ddp_spawn",
+    )
+    model = MultiProcessModel()
+    model.train_dataloader = None
+    model.val_dataloader = None
+    trainer.fit(model, train_dataloader=dataloader)
+
+
 def test_warning_with_iterable_dataset_and_len(tmpdir):
     """ Tests that a warning message is shown when an IterableDataset defines `__len__`. """
     model = BoringModel()