Context
Our current training pipeline relies on a standard PyTorch Dataset reading loose image files (e.g., COCO's 118,000 JPEGs) from NVMe/cloud storage. On high-core-count instances, increasing num_workers to saturate the GPU leads to severe OS lock contention and PCIe thread thrashing (e.g., 15% CPU utilization on 48 cores). The OS scheduler is overwhelmed by thousands of concurrent random file open() requests.
Proposed Solution
Introduce an alternative DataLoader pipeline using WebDataset.
The implementation should:
- Provide a utility script to pack the COCO dataset (images + JSON annotations) into ~100MB
.tar shards.
- Implement a
webdataset.WebLoader that streams these .tar shards sequentially.
- Map our existing Albumentations functions over the decoded WebDataset stream on the CPU workers.
- Yield the batched tensors to
pin_memory, handing off to our existing Kornia GPU pipeline.
Impact
- Eliminate I/O Bottlenecks: Converts thousands of slow, random IOPS into fast, sequential block reads.
- Zero Augmentation Rewrite: Perfectly compatible with our current Albumentations (CPU) and Kornia (GPU) architecture.
- Scalability: Makes the training code cloud-native, allowing us to eventually stream data directly from S3/GCS buckets without downloading the dataset locally first.
- System Stability: Drastically reduces OS thread thrashing and memory bloat associated with multi-processing loose file reads.
Context
Our current training pipeline relies on a standard PyTorch
Datasetreading loose image files (e.g., COCO's 118,000 JPEGs) from NVMe/cloud storage. On high-core-count instances, increasingnum_workersto saturate the GPU leads to severe OS lock contention and PCIe thread thrashing (e.g., 15% CPU utilization on 48 cores). The OS scheduler is overwhelmed by thousands of concurrent random fileopen()requests.Proposed Solution
Introduce an alternative DataLoader pipeline using
WebDataset.The implementation should:
.tarshards.webdataset.WebLoaderthat streams these.tarshards sequentially.pin_memory, handing off to our existing Kornia GPU pipeline.Impact