Search before asking
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?
Search before asking
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
WeightedMultiSourceBatchSamplerthat fixes the composition ofevery batch instead. With
batch_size=16and 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:
Design points I would suggest:
exactly to
batch_sizeand every source is represented wheneverbatch_sizeallows.when sources differ in size by orders of magnitude.
epoch_length="smallest"instead endsthe epoch when the smallest source has been seen once, so it is not repeated.
DistributedSamplershards samples, with theglobal 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
Samplerinrfdetr/datasets/, with no change toany existing training path. Users opt in by subclassing
RFDETRDataModuleand overridingtrain_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?