torch_ecg.utils.worker_init_fn#

torch_ecg.utils.worker_init_fn(worker_id: int) → None[source]#

worker_init_fn for torch.utils.data.DataLoader which reseeds DEFAULTS.RNG in each worker process.

PyTorch already reseeds the global random, torch and numpy.random states in each worker (with base_seed + worker_id, where base_seed derives from the generator of the data loader), but it knows nothing about numpy.random.Generator instances held by the user. Via fork (the default start method on Linux), each worker inherits a copy of the parent’s DEFAULTS.RNG, so all workers would produce identical random streams, e.g. for augmentations driven by DEFAULTS.RNG.

This function derives a distinct, reproducible seed from the worker-specific torch seed (torch.initial_seed()) and reseeds DEFAULTS.RNG (and the RNG_sample / RNG_randint partials) accordingly.

Usage

from torch.utils.data import DataLoader
from torch_ecg.utils import seed_everything, worker_init_fn

seed_everything(42)
dl = DataLoader(
    ds,
    batch_size=32,
    num_workers=4,
    worker_init_fn=worker_init_fn,
)
Parameters:

worker_id (int) – Id of the worker process, passed by torch.utils.data.DataLoader.

On this page