Table of Contents
Shortcuts

StratifiedBatchSampler#

class ignite.distributed.auto.StratifiedBatchSampler(labels, batch_size, strategy='oversample', num_replicas=None, rank=None, seed=0)[source]#

Batch sampler guaranteeing that every group is represented in every batch.

Unlike torch WeightedRandomSampler, which only makes rare groups more likely, this sampler puts a fixed number of samples of each group into each batch. This is required by methods where a missing group zeroes out a loss term, e.g. Group Distributionally Robust Optimization or contrastive learning.

It is a batch sampler and should be passed to batch_sampler of a dataloader, not to sampler:

batch_sampler = StratifiedBatchSampler(group_ids, batch_size=32)
dataloader = DataLoader(dataset, batch_sampler=batch_sampler)

@trainer.on(Events.EPOCH_STARTED)
def set_epoch(engine):
    batch_sampler.set_epoch(engine.state.epoch)
Parameters:

  • labels (Any) – group id of every dataset sample, as a sequence, numpy array or 1-D tensor. Groups are derived from the distinct values, which can be of any hashable type.

  • batch_size (int) – number of samples per batch. Should be larger or equal to the number of groups. Every yielded batch has exactly this size. If it is not divisible by the number of groups, the leftover slots are given to a rotating subset of the groups.

  • strategy (str) – if "oversample" (default), an epoch lasts until the largest group is exhausted and smaller groups are cycled, such that no sample is discarded. If "undersample", an epoch stops once the smallest group is exhausted and larger groups are truncated. Groups are reshuffled at every epoch, so the truncated part of a larger group is different at every epoch and no sample is permanently unused.

  • num_replicas (int | None) – number of processes participating in distributed training. Defaults to get_world_size().

  • rank (int | None) – rank of the current process within num_replicas. Defaults to get_rank().

  • seed (int) – random seed used to shuffle the groups. The seed of an epoch is seed + epoch.

Note

In a distributed configuration, groups are sharded across ranks individually, so that every batch of every rank still contains every group. Wrapping this sampler with DistributedProxySampler or torch DistributedSampler instead would shard the flat index list and could leave a rank without any sample of a rare group. Each group should therefore have at least num_replicas samples. Groups are truncated to a multiple of num_replicas, so that all ranks yield the same number of batches.

Note

As for torch DistributedSampler, set_epoch() should be called at the beginning of every epoch, otherwise the same batches are produced at every epoch.

New in version 0.6.0.

Methods

set_epoch

Set the epoch used to seed the shuffling of the groups.
set_epoch(epoch)[source]#

Set the epoch used to seed the shuffling of the groups.

Parameters:

epoch (int) – epoch number.

Return type:

None

×

Search Docs