Adding StratifiedBatchSampler for guaranteed subgroup representation - #3817
Open
avishkarsonni wants to merge 2 commits into
Open
Adding StratifiedBatchSampler for guaranteed subgroup representation#3817avishkarsonni wants to merge 2 commits into
avishkarsonni wants to merge 2 commits into
Conversation
- Adding this batch sampler as a feature so that sample of every grp is added to every batch for Group DRo and contrastive learning (methods where missing group zeroes a loss term) - 'oversample' is by default and 'undersample' to discard the smallest group In a distributed configuration, groups are sharded across ranks individually, so that every batch of every rank still contains every group Subclasses torch BatchSampler rather than Sampler, so that a DeterministicEngine can wrap it in ReproducibleBatchSampler.
TahaZahid05
reviewed
Jul 24, 2026
| return iter(indices) | ||
|
|
||
|
|
||
| class StratifiedBatchSampler(BatchSampler): |
Collaborator
There was a problem hiding this comment.
@vfdev-5 I am not sure if this is the correct place to put the sampler class. However, we really dont have a samplers/ folder either in ignite. Would it make sense to have ignite/samplers? like we have for metrics, handlers and so on.
Contributor
Author
There was a problem hiding this comment.
I think since DistributedProxySampler broke this rule years ago and creating new ignite/samplers directory will basically break the public path and create all new boilerplate, so I don't think it works well and it is for only 2-3 sampler as of now
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #3653
Description:
This PR adds
StratifiedBatchSampleras a guaranteed way to add fixed number of samples of every group in every batch for the methods where missing groups causes in omission of loss terms (Group DRO and Contrastive learning) which adds another method to existing methods likeWeightedRandomSampler,DistributedSampler,auto_dataloaderCheck list: