class

WeightedRandomSampler

extendsSampler

WeightedRandomSampler(weights: Sequence[float] | Tensor, num_samples: int, replacement: bool = True, generator: Generator | int | None = None)

source edit

Yield indices drawn proportionally to user-supplied weights.

Each index i is selected with probability proportional to weights[i]. Internally the weights are normalised by their sum so they need not form a true probability distribution on input.

Parameters

weightssequence of float or Tensor

Non-negative, finite weight per index. The effective probability of index i is

p_i = w_i / \sum_j w_j

num_samplesint

Number of indices to draw per epoch.

replacementbool= True

If True (default), draws are independent with replacement — the same index may appear multiple times. If False, each index is drawn at most once: the pool shrinks as indices are taken, so num_samples may not exceed len(weights) and is refused at construction if it does.

generatorlucid.Generator or int= None

Stream the draws come from; advances from epoch to epoch, so each epoch is a new draw and a fixed generator reproduces the sequence of epochs. None (default) draws from the default generator, so lucid.manual_seed reproduces the draws.

Notes

Each index $i$ is selected with probability

P(i) = \frac{w_i}{\sum_j w_j},

so the user-supplied weights need not be normalised. The classic use case is class-imbalance correction: set w_i = 1 / class_count[label_i] so under-represented classes are upsampled to roughly uniform frequency. Sampling with replacement is the default — it preserves the target marginal exactly and is the only fully consistent option when num_samples exceeds the number of nonzero-weight indices.

With replacement each draw inverts the cumulative weights by binary search, $O(n + k \log n)$ per epoch for k draws. Without replacement the indices are ranked by the exponential keys $\log(u_i) / w_i$ (Efraimidis & Spirakis, 2006) and the num_samples largest are kept, in the order one-at-a-time draws would pick them. When fewer than num_samples weights are positive, the draw stops after the last of them rather than handing out an index the caller weighted out.

Examples

>>> # 3 classes, counts [900, 90, 10]; upweight rare classes
>>> import lucid
>>> from lucid.utils.data import TensorDataset, WeightedRandomSampler
>>> labels = [0] * 900 + [1] * 90 + [2] * 10
>>> my_dataset = TensorDataset(lucid.randn(1000, 4), lucid.tensor(labels))
>>> weights = [1/900]*900 + [1/90]*90 + [1/10]*10
>>> sampler = WeightedRandomSampler(weights, num_samples=1000, generator=0)
>>> for idx in sampler:
...     x, y = my_dataset[idx]
>>> counts = [0, 0, 0]
>>> for idx in sampler:
...     counts[labels[idx]] += 1
>>> all(250 < c < 420 for c in counts)       # the classes now arrive evenly
True

Used by 1

lucid.utils.data

Constructors

dunder

init

→None

__init__(weights: Sequence[float] | Tensor, num_samples: int, replacement: bool = True, generator: Generator | int | None = None)

source edit

Store sampling configuration.

Parameters

weightssequence of float or Tensor

See class docstring.

num_samplesint

See class docstring.

replacementbool= True

See class docstring.

generatorlucid.Generator or int= None

See class docstring.

Raises

ValueError

If num_samples is not positive, exceeds len(weights) without replacement, or a weight is negative or not finite, or — with replacement — every weight is zero.

Dunder methods

dunder

iter

→Iterator[int]

__iter__()

source edit

Yield num_samples indices proportional to the weights.

With replacement=True each index comes from an inverse-CDF lookup (binary search over the running weight sums). With replacement=False the indices with the largest exponential keys are taken. See the class notes.

dunder

len

→int

__len__()

source edit

Return num_samples.

>>> # 3 classes, counts [900, 90, 10]; upweight rare classes >>> import lucid >>> from lucid.utils.data import TensorDataset, WeightedRandomSampler >>> labels = [0] * 900 + [1] * 90 + [2] * 10 >>> my_dataset = TensorDataset(lucid.randn(1000, 4), lucid.tensor(labels)) >>> weights = [1/900]*900 + [1/90]*90 + [1/10]*10 >>> sampler = WeightedRandomSampler(weights, num_samples=1000, generator=0) >>> for idx in sampler: ... x, y = my_dataset[idx] >>> counts = [0, 0, 0] >>> for idx in sampler: ... counts[labels[idx]] += 1 >>> all(250 < c < 420 for c in counts) # the classes now arrive evenly True