spdl.dataloader.PyTorchDataLoader¶
- class PyTorchDataLoader[source]¶
A PyTorch-style data loader that works on map-style dataset. Use
get_pytorch_dataloader()to instantiate this class. You can use this class as almost drop-in replacement of PyTorch’s DataLoader class.The architecture of data loader is different in following ways:
Only the dataset and the collate function are copied to the worker process. (Sampler and Generator are not copied)
The dataset is copied to worker processed via shared memory.
Sampler is executed in the main process and the resulting indices are passed to the worker processes.
Worker processes share the same input/output queues. (PyTorch creates a set of i/o queues for each worker process.)
Due to the way Dataset is defined, this class still has to copy the dataset to each worker process. So the memory consumption is not reduced. However, fast initialization and reduced inter-process communication makes this implementation faster than PyTorch DataLoader.
- Ivar:
dataset: The source dataset.
- Parameters:
dataset – Map-style dataset copied into each worker process.
shmem – Shared-memory buffer containing the serialized dataset.
sampler – Sampler executed in the main process.
fetch_fn – Function that resolves sampled indices in worker processes.
collate_fn – Function that collates samples in worker processes.
transfer_fn – Optional function that transfers a batch after collation.
mp_ctx – Multiprocessing context used to launch workers.
num_workers – Number of worker processes.
timeout – Maximum time to wait for a batch, or
Nonefor no timeout.buffer_size – Number of batches buffered by the pipeline.
output_order – Whether results are emitted in input or completion order.
worker_init_concurrency – Maximum number of workers that may deserialize the dataset concurrently.
Nonepreserves unbounded initialization.