feat(datasets): stream training input from WebDataset tar shards

Author: JESUSROYETH

.github/workflows/ci-tests-cpu.yml

@@ -56,7 +56,9 @@ jobs:
         # Install PyTorch CPU-only first (UV_TORCH_BACKEND=cpu works with 'uv pip').
         # [coreml] added on macOS only (coremltools is a macOS-only dependency).
         run: |
-          EXTRAS="train,augment,cli,visual"
+          # webdataset gates the streaming half of tests/datasets/test_webdataset_io.py behind
+          # importorskip; without it those tests silently skip and only the packer stays covered.
+          EXTRAS="train,augment,cli,visual,webdataset"
           if [ "${{ runner.os }}" = "macOS" ]; then
             EXTRAS="$EXTRAS,coreml"
           fi

CHANGELOG.md

@@ -8,6 +8,16 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),
 
 ### Added
 
+- Added `dataset_file="webdataset"`, an opt-in training input path that streams a split from pre-packed `.tar` shards instead of opening one file per image, together with `python -m rfdetr.datasets.webdataset_io` to pack a COCO split into ~100 MB shards and a small JSON index. Shards are written with the standard library, so packing needs no extra dependency; reading them installs with `pip install "rfdetr[webdataset]"`. Image bytes are copied verbatim, and a decoded sample goes through the same `ConvertCoco` conversion, the same `make_coco_transforms` CPU pipeline (Albumentations included), the same collate function and the same `pin_memory` hand-off to the Kornia GPU stage as the loose-file loaders — verified by comparing a packed split against the directory it came from, tensor by tensor, on 500 real COCO images: identical pixels, boxes, labels and label space. `coco`, `roboflow` and `yolo` are unaffected.
+
+    **Epoch sizing.** A streaming dataset has no index to sample from, so the training loader gives every worker a fixed sample count floored to a whole number of accumulation windows (`batch_size × grad_accum_steps`), the streaming counterpart of the padding `GradAccumAlignedDataset` gives the map-style loader: the reported length equals the batches produced, which is what `trainer.estimated_stepping_batches` — and through it the LR and drop schedules — needs, and no accumulation window is ever left partial. Validation, test and predict instead let every worker drain its shards once, so a split is scored exactly once and those loaders report no length. A packed `test` split is used when present, falling back to `val` with a log line otherwise. Training refuses to start when `world_size × num_workers` exceeds the shard count, and warns when an uneven shard split leaves the worst-served worker more than 5% short of the samples its epoch asks for — measured from each shard's real sample count when the packer recorded them, since shards are cut by byte size and a count-only estimate can be badly wrong when image sizes vary. Shard order and an in-stream reservoir buffer are both reshuffled each epoch, reseeded from a per-worker epoch counter so a `persistent_workers=True` loader (the DataModule's own default whenever `num_workers > 0`) reshuffles too instead of replaying its first epoch's order — a local shuffle, not the global permutation `shuffle=True` gives a map-style loader.
+
+    **Throughput, measured on one machine.** On a 32-vCPU instance with an NVMe-attached persistent disk, streaming COCO 2017 shards gave the same loader throughput as the loose files across six configurations (0.92x–1.03x; page cache cold and warm, 8 and 32 workers, default and Albumentations augmentation). That pipeline is not I/O-bound: `iowait` stayed under 1%, one fully page-cached configuration read 0 MB from disk at the same speed as its cold twin, and a read-only benchmark on the same disk sustained roughly 5.5x the images per second the loader consumed — with shards again at 0.99x loose, both arms converging on ~192 MB/s against a ~250 MB/s single-stream device ceiling. What the streaming path does improve there is construction: building the train and val datasets went from a median 19.9 s and 2,810 MB of resident memory to 0.012 s and 0.9 MB over three cold runs per arm, per process and so per DDP rank, because the 450 MB annotation file is never parsed. Packing should pay off where per-file access genuinely is the constraint — network filesystems, object-store mounts, directories of millions of small files — which this disk is not.
+
+    **Accuracy is not established either way.** Over five seeds of a 3-epoch, 2,000-image fine-tune at `num_workers=8`, the streaming median landed about 0.011 `mAP@50:95` below the map-style loader when the split was packed into 39 shards, and matched it (0.3685 against 0.3681) after re-packing the identical split into 290 shards — the shard-count skew above. The streaming arm ran 2–3x the seed-to-seed spread of the map-style one in both cases. Parity on a full-length run over a real dataset was not measured.
+
+    **Scope.** Detection and segmentation splits. Keypoint training rejects this format explicitly: its label space is inferred from a whole parsed COCO annotation file, which a shard index does not carry. `num_classes` must be set on the model, because the auto-detection the other formats use looks for `train/_annotations.coco.json` or `data.yaml`; class names are read from the shard index and need no such help. Shard paths resolve against the local `dataset_dir` — streaming shards straight from object storage is not wired up here. ([#1392](https://github.com/roboflow/rf-detr/issues/1392))
+
 - Added `TrainConfig.eval_batch_size`, decoupling the validation, test and predict dataloaders from the training micro-batch size. The three eval loaders previously always reused the resolved `batch_size`, so lowering it to fit an optimizer step also shrank evaluation batches. Evaluation runs under `no_grad`, which avoids autograd activation storage, but in-fit validation still shares device memory with the model and optimizer state and needs memory for its own forward outputs. The default `None` inherits `batch_size` exactly as before. Unlike `batch_size` it accepts no `"auto"`: an explicit `eval_batch_size` is never probed and stays usable on the `batch_size="auto"` path, while leaving it unset keeps the existing "auto was never resolved" error for eval loaders too. The training dataloader, including its `grad_accum_steps` alignment padding, is unaffected.
 
 - `deploy_to_roboflow()`'s `version` argument is now optional: when omitted, the highest existing dataset version of the target project is resolved automatically via the Roboflow API (falling back to version `1` for a project with no generated versions, where the Roboflow SDK then raises its usual "Version number 1 is not found."). Passing an explicit `version` behaves exactly as before, with no extra API call. ([#1116](https://github.com/roboflow/rf-detr/issues/1116))

docs/learn/train/dataset-formats.md

@@ -380,6 +380,105 @@ model.train(
 
 ---
 
+## WebDataset Shards (Sequential I/O)
+
+A COCO or YOLO split stored as loose files costs one `open()` per image per epoch. Raising `num_workers` to keep a GPU fed multiplies that into thousands of concurrent random opens — cheap against a local NVMe disk, expensive against network storage, an object-store mount, or any filesystem whose metadata operations are slow.
+
+`dataset_file="webdataset"` reads the same images from a handful of `.tar` shards instead. Each worker walks its own shards front to back, so an epoch becomes a set of large sequential reads and one `open()` per shard rather than one per image. Augmentation is unchanged: shards decode into the same `(image, target)` pairs the loose-file loader produces, go through the same CPU Albumentations/torchvision pipeline, and reach the GPU through the same `pin_memory` hand-off and Kornia stage.
+
+This is an alternative input path, not a replacement — `coco`, `roboflow` and `yolo` behave exactly as before.
+
+!!! note "It is not a general speed-up"
+
+    Whether packing helps depends on where your input pipeline is actually blocked. Measured on a 32-vCPU machine with an NVMe-attached persistent disk, on COCO 2017, throughput was the same either way — between 0.92x and 1.03x across six configurations. That pipeline was limited by decode and augmentation, not by file access: `iowait` stayed under 1%, and raw reads on the same disk sustained roughly 5.5x the images per second the loader consumed. Packing pays off where per-file access genuinely is the constraint — network filesystems, object-store mounts, directories with millions of small files. What it improves regardless is construction.
+
+### What it does change, everywhere
+
+Building the dataset stops parsing the split's annotation file. `dataset_file="coco"` loads `instances_train2017.json` into a `pycocotools` index before the first batch; the streaming path reads a small JSON listing shards, sample count and categories, and takes each image's annotations from the shard next to its pixels.
+
+On full COCO 2017 that is a median 19.9 s and 2,810 MB of resident memory before the first batch, against 0.012 s and 0.9 MB (three cold runs per arm; the spread is under 0.13 s and RSS is identical to the tenth of a MB). Both are per process, so per rank under DDP.
+
+### Packing a split
+
+Shards are written with the standard library, so packing needs no extra dependency:
+
+```bash
+python -m rfdetr.datasets.webdataset_io \
+    --image-dir /data/coco/train2017 \
+    --annotations /data/coco/annotations/instances_train2017.json \
+    --output-dir /data/coco-shards \
+    --split train
+```
+
+Repeat once per split, into the same `--output-dir`:
+
+```
+coco-shards/
+├── train-000000.tar        # ~100 MB each, override with --max-shard-mb
+├── train-000001.tar
+├── ...
+├── train-index.json        # shard list, sample count, categories
+├── val-000000.tar
+└── val-index.json
+```
+
+Image bytes are copied verbatim — no re-encode — so a packed split decodes to exactly the same pixels as the directory it came from. A `.json` member next to each image carries that image's `image_id`, `file_name` and annotation list, segmentation polygons included.
+
+`--category-ids` chooses the label space, and the index records the choice so the loader never has to guess:
+
+| Value             | Labels                                                       | Matches                   |
+| ----------------- | ------------------------------------------------------------ | ------------------------- |
+| `remap` (default) | contiguous `0..N-1`, unannotated grouping categories dropped | `dataset_file="roboflow"` |
+| `raw`             | the source `category_id` values                              | `dataset_file="coco"`     |
+
+### Training from shards
+
+```python
+from rfdetr import RFDETRSmall
+
+# num_classes is a model argument, not a train() one. 80 matches the `--category-ids remap`
+# default the packing command above used; see below for the `raw` formula.
+model = RFDETRSmall(num_classes=80)
+model.train(
+    dataset_dir="/data/coco-shards",
+    dataset_file="webdataset",
+    epochs=10,
+    batch_size=16,
+    num_workers=16,
+)
+```
+
+Install the reader with `pip install "rfdetr[webdataset]"`.
+
+`num_classes` has to be given explicitly: the auto-detection the other formats use looks for `train/_annotations.coco.json` or `data.yaml`, and a shard directory has neither, so the model keeps whatever count it was built with. Read the right number from `train-index.json` — under `raw`, that is the **highest `id` in `categories`, plus one** (COCO's own ids run 1-90 with gaps for its 80 categories, so `len(categories)` undercounts it — the same `max_obj_id + 1` convention `dataset_file="coco"` uses); under `remap`, it is simply the number of categories left after grouping nodes are dropped. Class *names* need no such help: they are read from the same index, indexed by the same labels the loader emits, so checkpoints and `predict()` label output are unaffected.
+
+### How an epoch is sized
+
+A streaming dataset has no index to sample from, so the two loaders size an epoch differently:
+
+- **Training** gives every worker a fixed number of samples, floored to a whole number of *accumulation windows* (`batch_size × grad_accum_steps`), so a partial window at the tail never fires the optimizer early. The epoch length is then exact, which is what the LR schedule needs — it is derived from `trainer.estimated_stepping_batches`. The cost is that a worker holding fewer shards than average repeats some of its own samples inside the epoch.
+- **Validation and test** let every worker drain its shards once, so the split is scored exactly once with nothing repeated or dropped. Those loaders report no length, so their progress bar shows no total.
+
+The test stage uses a packed `test` split when the directory has one, and falls back to `val` with a log line when it does not — a shard directory only has a split someone packed deliberately.
+
+Training raises a `ValueError` rather than starting if `world_size × num_workers` exceeds the shard count, because a worker left without a shard would silently shorten every epoch. Pack with a smaller `--max-shard-mb` to get more shards, or lower `num_workers`.
+
+### Pack plenty of shards
+
+Shards are split across workers by count, so an uneven split leaves the worst-served worker short of the samples a fixed epoch asks of it: it repeats some of its own while better-supplied workers leave some unseen. This is measurable. On a 2,000-image, 3-epoch fine-tune at `num_workers=8`, packing into 39 shards — the short worker 18% under the average — scored about 0.011 `mAP@50:95` below the loose-file loader; re-packing the identical split into 290 shards, 0.6% under, matched it. The loader logs a warning when that shortfall passes 5%, measured from each shard's real sample count (shards are cut by byte size, so a count-only estimate can be badly wrong when image sizes vary a lot within a split).
+
+Aim for a shard count that divides `world_size × num_workers`, or simply for many more shards than workers. `--max-shard-mb` is the knob.
+
+### Shuffling
+
+Shard visiting order is shuffled, and samples are shuffled again in a reservoir buffer as they stream. That is a local shuffle, not the global permutation `shuffle=True` gives a map-style loader: two samples in the same shard stay more likely to land in the same epoch region. Both are reseeded every epoch from PyTorch's own per-epoch worker seeding, so `seed_everything` governs them the same way it governs the rest of training.
+
+### Not covered
+
+Keypoint training rejects `dataset_file="webdataset"` with an explicit error. Its label space is inferred from a whole parsed COCO annotation file, which a shard index does not carry — use `coco`, `roboflow` or `yolo` for keypoints. Detection and segmentation splits are supported.
+
+---
+
 ## Converting Between Formats
 
 ### YOLO to COCO

pyproject.toml

@@ -64,6 +64,13 @@ augment = [
     "albumentations>=1.4.24,<3.0.0",  # optional custom CPU augmentations
     "kornia>=0.7,<1",                  # optional GPU-side augmentation
 ]
+webdataset = [
+    # Streaming tar-shard training input (dataset_file="webdataset"). Only the reader needs it: shards are packed
+    # with the standard library's tarfile, so `python -m rfdetr.datasets.webdataset_io` runs without this extra.
+    # Floored at 1.0, the series whose WebDataset(...) signature this integrates against: the loader
+    # passes workersplitter, nodesplitter, empty_check and seed explicitly.
+    "webdataset>=1.0,<2",
+]
 onnx = [
     "onnx>=1.16.0,<2.0",
     "onnxsim>=0.7.0",  # 0.7.0 ships wheels for cp310/311/312-abi3 (3.10-3.13) on linux x86_64+aarch64, win_amd64 and macOS arm64. The old <0.6.0 pin resolved to 0.5.0, which lacks cp311/cp313/aarch64 wheels, so pip built onnxsim's bundled onnxruntime/onnx from source and the install hung (#749).
@@ -399,6 +406,7 @@ overrides = [
     "coremltools", "coremltools.*",
     "roboflow", "roboflow.*",
     "rfdetr_plus", "rfdetr_plus.*",
+    "webdataset", "webdataset.*",
   ], ignore_missing_imports = true },
   # Optional-import None assignments in console.py: type: ignore[assignment, misc] is needed when
   # rich IS installed (full-venv mypy sees the assignment error) but flagged as unused by the

src/rfdetr/config.py

@@ -1060,7 +1060,9 @@ class TrainConfig(BaseConfig):
     keypoint_visible_loss_coef: float = 0
     keypoint_nll_loss_coef: float = 0
     keypoint_oks_sigmas: list[float] | None = None
-    dataset_file: Literal["coco", "o365", "roboflow", "yolo"] = "roboflow"
+    # "webdataset" streams pre-packed tar shards instead of loose image files; see
+    # rfdetr.datasets.webdataset_io for the packer and the sizing contract it imposes on the loaders.
+    dataset_file: Literal["coco", "o365", "roboflow", "yolo", "webdataset"] = "roboflow"
     square_resize_div_64: bool = True
     dataset_dir: PathLikeStr | None
     output_dir: PathLikeStr = "output"

src/rfdetr/datasets/__init__.py

@@ -26,6 +26,7 @@
 from rfdetr.datasets._keypoint_schema import infer_yolo_keypoint_schema as infer_yolo_keypoint_schema
 from rfdetr.datasets.coco import build_coco, build_roboflow_from_coco
 from rfdetr.datasets.o365 import build_o365
+from rfdetr.datasets.webdataset_io import build_webdataset
 from rfdetr.datasets.yolo import YoloDetection, build_roboflow_from_yolo
 
 
@@ -95,4 +96,6 @@ def build_dataset(image_set: str, args: Any, resolution: int) -> Dataset[Any]:
         return build_roboflow(image_set, args, resolution)
     if args.dataset_file == "yolo":
         return build_roboflow_from_yolo(image_set, args, resolution)
+    if args.dataset_file == "webdataset":
+        return build_webdataset(image_set, args, resolution)
     raise ValueError(f"dataset {args.dataset_file} not supported")