feat(datasets): add weighted multi-source batch sampler by Maryyyyyyyam142 · Pull Request #1287 · roboflow/rf-detr

Maryyyyyyyam142 · 2026-08-05T12:31:39Z

Description

Training on a mix of datasets through a plain ConcatDataset samples each source in proportion to
its size, so a large public dataset dominates every batch and a small hand-labelled set contributes
almost nothing to the gradient.

WeightedMultiSourceBatchSampler fixes the composition of every batch instead. With
batch_size=16 and weights [0.6, 0.3, 0.1], every batch holds exactly 10 / 5 / 1 samples from
the three sources.

from torch.utils.data import ConcatDataset, DataLoader
from rfdetr.datasets import WeightedMultiSourceBatchSampler

dataset = ConcatDataset([labeled, synthetic, public])
sampler = WeightedMultiSourceBatchSampler.from_concat_dataset(dataset, [0.6, 0.3, 0.1], batch_size=16)
loader = DataLoader(dataset, batch_sampler=sampler, collate_fn=collate_fn)

Implementation notes:

Slot allocation uses the largest-remainder (Hamilton) method on integer-scaled weights, so the
per-source counts always sum to batch_size and ties resolve deterministically rather than by
floating-point rounding error. Every source gets at least one slot when batch_size >= len(weights).
Recycling: a source that runs out mid-epoch is reshuffled and reused, which is what keeps the
ratio exact when sources differ in size by orders of magnitude. A warning is logged when a source
is recycled 10x or more per epoch, since that risks overfitting it. Only sources that actually
contribute samples are considered, so a starved source cannot mask an over-recycled contributor.
Epoch length is driven by the largest source by default; epoch_length="smallest" ends the
epoch when the smallest source has been seen once, and an integer selects a specific source.
DDP: each rank consumes a disjoint stride of the global batch stream, mirroring
DistributedSampler. Every rank builds every batch so the shared RNG stays aligned, and the
global batch count is truncated to a multiple of num_replicas so all ranks run the same number
of steps and none stalls in gradient all-reduce. Users must pass use_distributed_sampler=False
to the Trainer; this is documented.
Shuffling is seeded from (seed, epoch) via set_epoch(), matching DistributedSampler
semantics.

The sampler is opt-in and self-contained. No existing training path changes: RFDETRDataModule,
build_dataset, and the default dataloaders are untouched. Users opt in by subclassing
RFDETRDataModule and overriding train_dataloader(), as shown in the docs.

Type of change

New feature (non-breaking change which adds functionality)

How has this change been tested, please provide a testcase or example of how you tested the change?

tests/datasets/test_multi_source.py adds 49 tests grouped in eight classes:

TestComputeSourceBatchSizes — parametrized allocation cases (documented example, even split,
remainders, tiny weights, batch_size smaller than the source count), the sum invariant, weight
normalisation, reclaiming slots when the one-per-source guarantee over-allocates, and rejection of
zero weights, empty weights, weights below the representable precision, and a non-positive
batch_size.
TestSamplerValidation — length mismatch, empty source, no sources at all, non-positive
num_replicas, rank outside the world, invalid epoch_length.
TestBatchComposition — every batch matches the requested ratio and full batch_size, indices
stay inside their source, small sources are recycled while the driving source is not repeated.
TestEpochLength — largest / smallest / explicit index, __len__ matches the number of
yielded batches, drop_last controls whether the trailing partial batch of the driving source is
kept (and that batch is still full batch_size), tiny datasets still yield one batch.
TestRecyclingWarning — the over-recycling warning names the most-recycled contributing source,
a starved source (zero slots per batch) neither suppresses the warning nor gets reported, the
reported recycling factor is correct, balanced sources warn nothing, and an epoch_length that
points at a starved source falls back to a contributing one.
TestShufflingDeterminism — same epoch reproduces, different epoch and different seed reshuffle,
shuffle=False is sequential.
TestDistributedSharding — equal batch counts per rank, disjoint batches across ranks, and the
interleaved 2-rank stream reproduces the single-process stream exactly.
TestDataLoaderIntegration — end-to-end use as a DataLoader batch sampler and the
from_concat_dataset constructor.

Verified locally:

pytest tests/datasets -n 2 -m "not gpu" → 572 passed, 70 skipped.
coverage run --include='*/rfdetr/datasets/multi_source.py' -m pytest tests/datasets/test_multi_source.py
→ 143 statements, 0 missed, 100% coverage of the new module.
pre-commit run --all-files → all hooks pass except mypy, which reports the same pre-existing
errors with and without this change (all import-stub and subclass errors in files this PR does not
touch); none mention the new module or its tests.
mkdocs build → succeeds, both new documentation sections render.

Any specific deployment considerations

None. The feature is additive, opt-in, and pure Python over torch.utils.data; no new dependency
and no change to default behaviour.

Training on a mix of datasets through a plain ConcatDataset samples each source in proportion to its size, so a large public dataset dominates every batch and a small hand-labelled set contributes almost nothing to the gradient. WeightedMultiSourceBatchSampler fixes the composition of every batch instead: batch slots are allocated across sources with the largest-remainder method, so the counts sum exactly to batch_size and every source is represented whenever batch_size allows. A source that runs out mid-epoch is reshuffled and reused, which is what keeps the ratio exact when sources differ in size by orders of magnitude. Epoch length is driven by the largest source by default, or by the smallest one to avoid repeating it. Batches are sharded across DDP ranks the way DistributedSampler shards samples, with the global batch count truncated to a multiple of the world size so no rank stalls in gradient all-reduce. The sampler is opt-in and self-contained; no existing training path changes.

CLAassistant · 2026-08-05T12:31:54Z

All committers have signed the CLA.

Maryyyyyyyam142 · 2026-08-05T12:41:52Z

I have read the CLA Document and I sign the CLA.

Copilot

Pull request overview

This PR introduces a new opt-in PyTorch batch sampler to support training on multiple concatenated datasets with a fixed per-batch source ratio, addressing the common issue where ConcatDataset sampling is dominated by the largest dataset.

Changes:

Added WeightedMultiSourceBatchSampler and compute_source_batch_sizes to enforce deterministic per-batch source composition, with optional DDP-aware batch stream sharding.
Added a comprehensive test suite covering allocation, validation, epoch semantics, determinism, DDP sharding, and DataLoader integration.
Documented usage in training customization and reference docs, and added a changelog entry.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
`tests/datasets/test_multi_source.py`	New tests validating weighted multi-source batch composition, epoch length behavior, determinism, and DDP sharding semantics.
`src/rfdetr/datasets/multi_source.py`	New implementation of weighted multi-source batch sampling + allocation helper.
`src/rfdetr/datasets/__init__.py`	Re-exported the sampler and allocation helper from `rfdetr.datasets`.
`docs/reference/training.md`	Added API reference entries for the sampler and allocation helper.
`docs/learn/train/customization.md`	Added a runnable example showing how to opt in via a custom `RFDETRDataModule`.
`CHANGELOG.md`	Added an `[Unreleased] / Added` entry describing the new feature.

Borda · 2026-08-05T21:46:31Z

I have read the CLA Document and I sign the CLA.

@Maryyyyyyyam142, could you pls share a screenshot... 🦝

codecov · 2026-08-05T21:48:32Z

Codecov Report

❌ Patch coverage is 99.01478% with 2 lines in your changes missing coverage. Please review.
✅ Project coverage is 85%. Comparing base (ec27026) to head (da125ed).
⚠️ Report is 2 commits behind head on develop.

Additional details and impacted files

@@           Coverage Diff            @@
##           develop   #1287    +/-   ##
========================================
  Coverage       85%     85%            
========================================
  Files          111     112     +1     
  Lines        13918   14120   +202     
========================================
+ Hits         11809   12009   +200     
- Misses        2109    2111     +2

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

_warn_on_source_imbalance() picked the smallest source by size. When batch_size is below the number of sources, the smallest source can be one that receives zero slots per batch, and a starved source is never recycled so its pass count is 0. That suppressed the warning entirely even when a different, non-starved source was being recycled heavily. Recycling is now computed only over sources that actually contribute samples, and the warning reports the most-recycled of them. With sources of 1000/20/50/5 samples at weights 0.5/0.3/0.15/0.05 and batch_size=3, sources 2 and 3 are starved; the warning previously stayed silent about source 1 being repeated 25x per epoch and now names it. Also restores the `shuffle` entry in the class Args docstring, dropped by the previous commit while expanding the `drop_last` entry.

…pler

Adds tests for the validation paths (empty weights, sub-precision weights, missing sources, non-positive num_replicas), the slot-reclaim early exit, and drop_last behaviour on a trailing partial batch. The module is now at 100% statement coverage.

Maryyyyyyyam142 · 2026-08-09T13:34:03Z

I have read the CLA Document and I sign the CLA.

Maryyyyyyyam142 · 2026-08-09T13:41:54Z

Pushed a follow-up that brings src/rfdetr/datasets/multi_source.py to 100% statement coverage
(143 statements, 0 missed), so the 93.2% patch-coverage report above is stale. The added tests
cover the validation paths (empty weights, sub-precision weights, no sources, non-positive
num_replicas), the slot-reclaim early exit in compute_source_batch_sizes, and drop_last on a
trailing partial batch. The branch is also up to date with develop now, which clears the
"8 commits behind head" note.

@Borda the workflows on the latest push are sitting in action_required — could you approve the
run so Codecov can re-report? I've also updated the PR description; the test count is 49 across
eight classes now.

…pler

GitHub's Update-branch conflict resolution left the WeightedMultiSourceBatchSampler bullet under the released 1.9.2 section. Move it to Unreleased ### Added.

[resolve group] PR roboflow#1287 — items 6, 8, 9, 10, 11 - compute_source_batch_sizes could return counts summing above batch_size (single-pass reclaim loop under-recovered slots); now multi-pass with a sum(counts) == batch_size post-condition. - Warn when a source's realized batch fraction diverges materially from its requested weight (min-1 clamp could invert the ratio silently). - Document the cross-rank identical-ConcatDataset-layout precondition that DDP correctness depends on. - Lead the class docstring with the drop_last/torch-convention divergence instead of burying it in the arg entry. - Vectorize _shuffled_indices (34x faster, bit-identical output). --- Co-authored-by: OpenAI Codex <codex@openai.com>

[resolve group] PR roboflow#1287 — items 12, 13 Docs example for WeightedMultiSourceBatchSampler overrode train_dataloader() wholesale, silently dropping DataLoader invariants (num_workers, worker_init_fn, etc.) and reaching into private _dataset_train/_collate_fn attributes. Add a public build_train_sampler() override point on RFDETRDataModule, consulted by train_dataloader() before its default sampling logic; update the docs example to use it instead.

[resolve gate] QA gate flagged zero test coverage for train_dataloader()'s non-None build_train_sampler() branch added in 3dbe1ce — a regression there would go undetected by CI. Add coverage for: default None passthrough, custom batch_sampler passed through as-is, DataLoader kwargs (collate_fn/num_workers/ pin_memory) preserved, and GradAccumAlignedDataset wrapping skipped.

Auto-detect the DDP layout: num_replicas and rank default to None and resolve through get_world_size()/get_rank(), so a distributed run shards the batch stream instead of replaying the identical stream on every rank; an explicit num_replicas=1 under a live world size above 1 now warns. Reject non-finite weights: the positivity guard also requires isfinite(), so an inf weight raises the documented ValueError instead of an OverflowError from round(). Reject sub-precision weights per index: a weight that scales to zero raises the "too small to be represented at 1e-6 precision" error naming the offending entries, instead of being clamped up to a full guaranteed slot. This is a superset of the former all-zero-total check, which is removed. Skip the per-epoch shuffle of starved sources (zero slots per batch) in __iter__; those streams were built in full and never read. Extract _validate_and_scale_weights() and _rebalance_counts() out of compute_source_batch_sizes(), which had grown to cyclomatic complexity 14 and 13 branches. Guard the Lightning DDP misuse in train_dataloader(): a custom batch sampler under a distributed strategy with use_distributed_sampler=True now raises a RuntimeError naming Trainer(use_distributed_sampler=False), instead of surfacing Lightning's bare TypeError about missing __init__ arguments. Lightning-internal flags are read defensively so an unknown version cannot raise a false alarm.

…LOG PR link - customization.md: store the WeightedMultiSourceBatchSampler built in build_train_sampler() on self._multi_source_sampler so the documented on_train_epoch_start().set_epoch() call has a concrete handle to reach, since PTL only auto-wires set_epoch for samplers exposed as dataloader.sampler / dataloader.batch_sampler.sampler. - customization.md: pass self._resolve_batch_size() instead of train_config.batch_size to from_concat_dataset(), since batch_size may still be the literal "auto" on the custom-DataModule/PTL-CLI path (only RFDETR.train() resolves it), and _resolve_batch_size() raises a clear RuntimeError instead of a cryptic int() ValueError from the sampler. - CHANGELOG.md: append the PR link to the WeightedMultiSourceBatchSampler Added entry, matching neighboring PR-referencing entries.

…poch_length, negative weights - roboflow#22: parametrized sweep asserting min(counts) >= 1 across skewed weight sets once batch_size >= len(weights) - roboflow#23: DDP tests for a starved driving source, covering both the clamp-to-1 len(sampler) path and the truncate-to-a-multiple-of- num_replicas path, with num_replicas>1 in both - roboflow#26: Examples doctests for _source_of/_batch_composition, and doctest +SKIP (caplog-dependent) for _capture_warnings/ _recycling_warnings/_ratio_warnings - roboflow#28: test asserting rfdetr.datasets re-exports resolve to the same objects as rfdetr.datasets.multi_source - roboflow#29: parametrize test_rejects_unknown_epoch_length with bool True/False cases (bool is an int subclass) - roboflow#30: parametrize weight rejection with negative/NaN/infinite weights, renamed to test_rejects_non_positive_weight

…p guard QA gate flagged two untested paths in the DDP-hardening commit: the live-process-group auto-detect/warning branch of _resolve_distributed_layout (only the no-process-group and fully-explicit paths had doctests), and _check_custom_sampler_owns_ddp's RuntimeError plus its wiring into train_dataloader(), both new in this batch.

[resolve group] PR roboflow#1287 — items 6, 8, 9, 10, 11 - compute_source_batch_sizes could return counts summing above batch_size (single-pass reclaim loop under-recovered slots); now multi-pass with a sum(counts) == batch_size post-condition. - Warn when a source's realized batch fraction diverges materially from its requested weight (min-1 clamp could invert the ratio silently). - Document the cross-rank identical-ConcatDataset-layout precondition that DDP correctness depends on. - Lead the class docstring with the drop_last/torch-convention divergence instead of burying it in the arg entry. - Vectorize _shuffled_indices (34x faster, bit-identical output). --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com> Co-authored-by: OpenAI Codex <codex@openai.com>

[resolve group] PR roboflow#1287 — items 12, 13 Docs example for WeightedMultiSourceBatchSampler overrode train_dataloader() wholesale, silently dropping DataLoader invariants (num_workers, worker_init_fn, etc.) and reaching into private _dataset_train/_collate_fn attributes. Add a public build_train_sampler() override point on RFDETRDataModule, consulted by train_dataloader() before its default sampling logic; update the docs example to use it instead. --- Co-authored-by: claude[bot] <209825114+claude[bot]@users.noreply.github.com>

Borda

@Maryyyyyyyam142, nice work! Let's add a simple cookbook about this... we can combine a few datasets from Universe

…pler

…ampler Add a short cookbook that downloads three Roboflow Universe detection exports, unifies their class names, and trains with a fixed 4/3/1 per-batch source ratio.

Maryyyyyyyam142 · 2026-08-17T12:31:06Z

@Borda added a short cookbook that mixes three Universe detection datasets
(hard-hat PPE, traffic, football) with WeightedMultiSourceBatchSampler.

Every batch is 4 / 3 / 1 from those sources regardless of dataset size — the
notebook prints that composition before the short training run. Class names
are unified across the three COCO exports so the head sees one label space.

Notebook: docs/cookbooks/multi-source-batch-sampler.ipynb

The build_train_sampler example was a doctest with every line commented out, so pytest-doctestplus never ran it. Use a Sphinx code-block instead, matching the rest of the package for snippets that are not executable in isolation.

…pler

…pler Keep both cookbook cards, combine the train_dataloader docstring (PackedTargets + DDP RuntimeError), and keep sampler tests alongside develop's pack_targets collate tests.

…pler

Maryyyyyyyam142 · 2026-08-31T07:53:14Z

@Borda friendly ping — both points from your Aug 13 review are addressed, and it's been quiet a couple of weeks, so flagging in case this slipped off the radar.

Since that review:

Cookbook added — docs/cookbooks/multi-source-batch-sampler.ipynb, mixing three Universe detection datasets (hard-hat PPE, traffic, football), registered in cards.yaml. It prints the 4/3/1 per-batch composition before a short training run.
Commented-out doctest replaced with a normal .. code-block:: python example on build_train_sampler, so doctestplus doesn't try to execute a partial class snippet.
Branch synced with develop (Aug 26) — no conflicts, up to date.

One thing I can't do from my side: the workflows on 320aef2 are all sitting at "action_required", so CPU/GPU/docs/mypy have never actually run. CLA, pre-commit.ci and Socket are green, but the test suites need a maintainer to approve the run. If you can kick those off, that's the last signal you'd need to review against.

Happy to rebase, split it up, or trim scope if any of that would help it land.

Borda · 2026-08-31T09:08:48Z

Hi, apologies for the delay, I was off last week and plan for checking it later this week...

Maryyyyyyyam142 requested review from Borda, SkalskiP, isaacrob and probicheaux as code owners August 5, 2026 12:31

Borda requested a lite review from Copilot August 5, 2026 18:28

Copilot started reviewing on behalf of Borda August 5, 2026 18:28 View session

Copilot AI reviewed Aug 5, 2026

View reviewed changes

Comment thread src/rfdetr/datasets/multi_source.py Outdated

Comment thread src/rfdetr/datasets/multi_source.py

Maryyyyyyyam142 and others added 5 commits August 6, 2026 08:29

Potential fix for pull request finding

86108cb

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

d1002c0

…pler

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

c1e5ab5

…pler

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

57fdcfc

…pler

github-actions Bot added the has conflicts label Aug 11, 2026

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

b1dbfa3

…pler

github-actions Bot removed the has conflicts label Aug 12, 2026

Maryyyyyyyam142 and others added 7 commits August 12, 2026 20:09

docs(changelog): move sampler entry to Unreleased Added

d058dce

GitHub's Update-branch conflict resolution left the WeightedMultiSourceBatchSampler bullet under the released 1.9.2 section. Move it to Unreleased ### Added.

fix(pre-commit): 🎨 auto format pre-commit hooks

d042283

Borda added 2 commits August 13, 2026 09:14

Borda reviewed Aug 13, 2026

View reviewed changes

Comment thread src/rfdetr/training/module_data.py Outdated

Borda previously approved these changes Aug 13, 2026

View reviewed changes

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

84dc864

…pler

Borda added the enhancement New feature or request label Aug 13, 2026

Maryyyyyyyam142 and others added 2 commits August 17, 2026 15:14

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

67b8e2f

…pler

fix(pre-commit): 🎨 auto format pre-commit hooks

c481699

pre-commit-ci Bot dismissed Borda’s stale review via 78e3f58 August 17, 2026 12:15

Maryyyyyyyam142 and others added 2 commits August 17, 2026 15:27

docs(cookbooks): mix Universe datasets with WeightedMultiSourceBatchS…

4bffc85

…ampler Add a short cookbook that downloads three Roboflow Universe detection exports, unifies their class names, and trains with a fixed 4/3/1 per-batch source ratio.

fix(pre-commit): 🎨 auto format pre-commit hooks

7351c84

Maryyyyyyyam142 requested a review from Borda August 18, 2026 05:54

Maryyyyyyyam142 added 2 commits August 18, 2026 19:12

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

efc50a9

…pler

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

aaf90e4

…pler

github-actions Bot added the has conflicts label Aug 24, 2026

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

ebec8c3

…pler Keep both cookbook cards, combine the train_dataloader docstring (PackedTargets + DDP RuntimeError), and keep sampler tests alongside develop's pack_targets collate tests.

github-actions Bot removed the has conflicts label Aug 25, 2026

Merge branch 'develop' into feat/1286-weighted_multi_source_batch_sam…

8878897

…pler

Maryyyyyyyam142 force-pushed the feat/1286-weighted_multi_source_batch_sampler branch from 320aef2 to 8878897 Compare August 31, 2026 12:34

Navigation Menu

Uh oh!

feat(datasets): add weighted multi-source batch sampler - #1287

feat(datasets): add weighted multi-source batch sampler#1287
Maryyyyyyyam142 wants to merge 27 commits into
roboflow:developfrom
Maryyyyyyyam142:feat/1286-weighted_multi_source_batch_sampler

Maryyyyyyyam142 commented Aug 5, 2026 •

edited

Loading

Uh oh!

CLAassistant commented Aug 5, 2026 •

edited

Loading

Uh oh!

Maryyyyyyyam142 commented Aug 5, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

Borda commented Aug 5, 2026

Uh oh!

codecov Bot commented Aug 5, 2026 •

edited

Loading

Uh oh!

Maryyyyyyyam142 commented Aug 9, 2026

Uh oh!

Maryyyyyyyam142 commented Aug 9, 2026

Uh oh!

Uh oh!

Borda left a comment •

edited

Loading

Uh oh!

Maryyyyyyyam142 commented Aug 17, 2026

Uh oh!

Maryyyyyyyam142 commented Aug 31, 2026

Uh oh!

Borda commented Aug 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Uh oh!

Conversation

Maryyyyyyyam142 commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Type of change

How has this change been tested, please provide a testcase or example of how you tested the change?

Any specific deployment considerations

Uh oh!

CLAassistant commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Maryyyyyyyam142 commented Aug 5, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Borda commented Aug 5, 2026

Uh oh!

codecov Bot commented Aug 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Maryyyyyyyam142 commented Aug 9, 2026

Uh oh!

Maryyyyyyyam142 commented Aug 9, 2026

Uh oh!

Uh oh!

Borda left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Maryyyyyyyam142 commented Aug 17, 2026

Uh oh!

Maryyyyyyyam142 commented Aug 31, 2026

Uh oh!

Borda commented Aug 31, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Maryyyyyyyam142 commented Aug 5, 2026 •

edited

Loading

CLAassistant commented Aug 5, 2026 •

edited

Loading

codecov Bot commented Aug 5, 2026 •

edited

Loading

Borda left a comment •

edited

Loading