DirichletPartitioner#

Dirichlet dataset partitioning: per-class label-skew split.

Reference:: Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. “Measuring the Effects of Non-Identical Data Distribution for Federated Visual Classification.” In NeurIPS Workshop on Federated Learning (2019).

class krum.primitives.data_partitioners.dirichlet.DirichletPartitioner[source]#

Bases: DataPartitioner

Dirichlet partitioner: per-class label-skew split, from near-IID to extreme imbalance.

For each class \(k\), draws a proportion vector \(p_k \sim \mathrm{Dirichlet}(\alpha, \dots, \alpha)\) of dimension \(n\) (one entry per worker, summing to 1), then gives worker \(w\) a \(p_{k,w}\) fraction of class \(k\)’s samples. Every sample is assigned to exactly one worker (no remainder is dropped, unlike IidPartitioner).

\(\alpha\) controls the skew: as \(\alpha \to \infty\), every \(p_k\) collapses to \((1/n, \dots, 1/n)\) (near-IID); as \(\alpha \to 0\), every \(p_k\) collapses to a one-hot vector (each class goes almost entirely to a single worker).

A worker can legitimately end up with zero samples of a class, or (for small enough \(\alpha\) and small \(n\)) even zero samples overall — an intentional consequence of extreme skew. Such a worker gets an empty (but valid) dataset. An empty dataset is handled the same way: every worker gets an empty dataset.

classmethod partition(dataset: Dataset[Any], /, *, n: int, alpha: float, seed: int = 42, **specialized: Any) → list[Subset[Any]][source]#

Split dataset across n workers via per-class Dirichlet skew.

Parameters:

dataset – Full dataset to partition across workers. Labels are read from dataset.targets when available (as for the torchvision datasets), otherwise by indexing every sample.
n – Number of workers to split the dataset across.
alpha – Concentration parameter of the per-class \(\mathrm{Dirichlet}(\alpha, \dots, \alpha)\) draw. Smaller values produce more extreme label skew.
seed – Random seed for the per-class Dirichlet draws and the within-class shuffle.
**specialized – Additional keyword arguments (unused).

Returns:

List of ``n`` datasets, one per worker.

Raises:

ValueError – If n < 1 or alpha <= 0.