Skip to content

Implement WebDataset pipeline for sequential I/O #1392

Description

@Borda

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:

  1. Provide a utility script to pack the COCO dataset (images + JSON annotations) into ~100MB .tar shards.
  2. Implement a webdataset.WebLoader that streams these .tar shards sequentially.
  3. Map our existing Albumentations functions over the decoded WebDataset stream on the CPU workers.
  4. 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.

Metadata

Metadata

Assignees

Labels

enhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions