Weighted multi-source batch sampling: fixed per-source ratio in every training batch

Search before asking

• I have searched the RF-DETR issues and found no similar feature requests.

Description

Training on a mix of datasets (hand-labelled, synthetic, public) currently means concatenating them, which samples each source in proportion to its size. A 50k-image public set and a 2k-image hand-labelled set produce batches that are ~96% public data, so the small high-quality source contributes almost nothing to the gradient. Re-weighting after the fact (loss weighting, oversampling on disk) is either indirect or duplicates data.

I would like to propose a WeightedMultiSourceBatchSampler that fixes the composition of every batch instead. With batch_size=16 and weights [0.6, 0.3, 0.1] over three sources, each batch holds exactly 10 / 5 / 1 samples from them, regardless of the relative source sizes.

Sketch of the API:

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)

Design points I would suggest:

• Slots are allocated with the largest-remainder (Hamilton) method so the per-source 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; epoch_length="smallest" instead ends the epoch when the smallest source has been seen once, so it is not repeated.
• Batches are sharded across DDP ranks the same way DistributedSampler shards samples, with the global batch count truncated to a multiple of the world size so no rank stalls in all-reduce.

It would be additive and opt-in: a standalone Sampler in rfdetr/datasets/, with no change to any existing training path. Users opt in by subclassing RFDETRDataModule and overriding train_dataloader().

Use case

Anyone fine-tuning RF-DETR on a small domain-specific dataset while mixing in synthetic or public data to avoid overfitting. It is also useful for curriculum-style ratios (start heavy on synthetic, shift toward real data) and for keeping a rare-class source present in every batch.

I have been running this approach on an RF-DETR training pipeline and would like to contribute it upstream in a form that fits the current Lightning-based codebase.

Additional

Happy to adjust the placement (rfdetr/datasets/ vs rfdetr/training/) and the API surface before implementing. I have a working implementation with unit tests and docs ready to open as a PR once the approach is approved.

Are you willing to submit a PR?

• Yes I'd like to help by submitting a PR!