SR-Forge

Loader

DataLoaderFactory

A factory for creating PyTorch DataLoader instances.

This factory inspects the dataset type (Entry or GraphEntry) and the training environment (single-GPU or multi-GPU) to select the most appropriate DataLoader from PyTorch or PyTorch Geometric.

It also owns the operational worker policy, so training scripts don't have to (every default below is overridable via **kwargs):

• multiprocessing_context='forkserver' (POSIX, when workers > 0) — plain fork ed workers inherit whatever the parent process carries at the fork instant (CUDA state, NCCL watchdogs, tracker/OpenCV threads, held locks) and intermittently deadlock in bootstrap: the silent "stuck at batch 0" hang. Forkserver workers are born from a clean helper process, immune by construction — and safe regardless of where in the script the loader is created.
• persistent_workers=True (when workers > 0) — the worker pool is started once and reused every epoch instead of being respawned.
• pin_memory is dropped automatically for CPU devices (PyTorch raises otherwise), so callers can pass it unconditionally.
• DEBUG_HANG=1 injects a stack-dumping worker_init_fn.

__init__(dataset: Dataset, batch_size: int = 32, shuffle: bool = True, device: Union[str, list] = 'cuda:0', **kwargs)

Initializes the DataLoaderFactory.

Parameters:

Name	Type	Description	Default
dataset	Dataset	The dataset to be loaded.	required
batch_size	int	The number of samples per batch. Defaults to 32.	32
shuffle	bool	Whether to shuffle the data at every epoch. Defaults to True.	True
device	Union[str, list]	The device(s) for training. If a list of device IDs is provided, multi-GPU loaders are configured. Defaults to 'cuda:0'.	'cuda:0'
**kwargs	Any	Additional keyword arguments to be passed to the underlying DataLoader constructor. These take precedence over the factory's worker-policy defaults (see class docstring).	{}

get_loader() -> DataLoader

Builds and returns the appropriate DataLoader.

Selects a loader based on the dataset's item type and device settings: - torch.utils.data.DataLoader: For Entry types on single/multi-GPU. - torch_geometric.loader.DataLoader: For GraphEntry on single-GPU. - torch_geometric.loader.DataListLoader: For GraphEntry on multi-GPU.

Under a distributed strategy (DDP), the active :class:~srforge.distributed.strategy.TrainingStrategy supplies a rank-aware sampler that shards the dataset across processes; shuffle is then delegated to that sampler (PyTorch forbids passing both).

Returns:

Name	Type	Description
DataLoader	DataLoader	The configured DataLoader instance.

Raises:

Type	Description
TypeError	If the dataset's item type is not supported.