Note about modification to PyTorch’s default dataloader

This section describes modifications to PyTorch’s default DataLoader used in SkiNet to prevent spawning new processes at the beginning of each epoch.

A Jupyter notebook with examples is available at: SkiNet/ML/dataloaders/examples/RepeatDataloaders.ipynb

Motivation

Spawning new worker processes at the beginning of each epoch in PyTorch’s DataLoader (when num_workers > 0 and persistent_workers=False) causes:

  • Overhead of initialising the dataset in each worker (by deserialising from the main process)

  • Memory overhead as each worker requires its own memory space

  • Any worker-cached state is lost when workers are shut down

  • Copy-on-write behaviour during fork increases memory if the dataset contains Python lists or dicts

  • Potential deadlocks and race conditions with shared resources

Dataloaders and subprocesses

When iterating over a DataLoader:

  1. Dataset initialisation runs in the main process (Dataset.__init__).

  2. DataLoader initialisation runs in the main process (DataLoader.__init__).

  3. Prefetching & Queues: At the start of epoch 0, PyTorch spawns num_workers separate subprocesses. Dataset indices are put in worker input queues (up to prefetch_factor × num_workers indices ahead). Each worker calls __getitem__ asynchronously; the main process collects results.

At the start of each epoch, for batch in loader calls iter(loader), which calls __iter__():

  • If persistent_workers=True and num_workers > 0 and no iterator exists yet, _get_iterator() creates a new iterator.

  • If an iterator exists (epoch > 0), its state is reset via _reset() — workers are reused.

  • If persistent_workers=False and num_workers > 0, _get_iterator() is called every epoch, creating new workers each time.

When num_workers > 0, _get_iterator() returns _MultiProcessingDataLoaderIter(self). New workers are created in _MultiProcessingDataLoaderIter.__init__(). Workers are normally shut down when the sampler is exhausted at epoch end.

How SkiNet prevents new worker spawning: RepeatDataLoader

SkiNet.ML.dataloaders.dataloaders.RepeatDataLoader subclasses torch.utils.data.DataLoader and relies on a single mechanism — persistent workers. This replaces an earlier _RepeatSampler/iterator-override implementation (the previous infinite-sampler approach is gone).

In RepeatDataLoader.__init__:

  • persistent_workers defaults to num_workers > 0 (set only if the caller did not pass it). With num_workers=0 the flag is invalid and PyTorch raises, so it is left False. Persistent workers survive across epochs: PyTorch’s _MultiProcessingDataLoaderIter is created once and _reset() between epochs instead of being torn down and respawned.

  • worker_init_fn defaults to default_worker_init_fn, which sets cv2.setNumThreads(0) and seeds numpy/random per worker from torch.initial_seed() % 2**32 for reproducibility.

  • collate_fn defaults to collate_preserving_specs (see below).

  • The dataset must be sized (__len__), else TypeError is raised.

  • max_num_to_repeat is accepted for backward compatibility but ignored — Lightning’s Trainer controls epoch iteration via max_epochs.

Usage with ISIC 2017

create_segmentation_dataloaders(main_config) is the segmentation entry point: it calls create_segmentation_datasets_from_config then create_dataloaders_from_datasets, which builds the RepeatDataLoader instances for the train, val, and test splits:

from SkiNet.ML.configs.load_config_from_yaml import load_config_from_yaml
from SkiNet.ML.dataloaders.create_dataloaders import create_dataloaders_from_datasets
from SkiNet.ML.datasets.dataset_factory import create_segmentation_datasets_from_config

cfg = load_config_from_yaml("main_config.yaml")
datasets = create_segmentation_datasets_from_config(cfg)
loaders = create_dataloaders_from_datasets(datasets, cfg.trainconfig)

# loaders.train, loaders.val, loaders.test are RepeatDataLoader instances

Key dataloader settings from TrainConfig that affect worker behaviour:

Field

Effect

num_workers

Auto-set to os.cpu_count() for single-GPU; divided by device count for DDP

pin_memory

Auto-set True on CUDA, False on MPS/CPU

prefetch_factor

Batches pre-loaded per worker; None when num_workers=0

cache_in_ram

Eliminates disk I/O in workers when True (recommended for small datasets)

Custom collate function

RepeatDataLoader uses collate_preserving_specs as its collate_fn. This preserves the specs field (sample metadata) as a Python list rather than attempting to stack it into a tensor, since metadata values are heterogeneous strings.