fedlib.datasets.DirichletPartitioner

class DirichletPartitioner(num_clients: int = 4, random_seed: int = 123, client_id_generator: Callable[[], Iterator] = None, alpha: float = 1.0, same_proportions: bool = True)

Bases: DatasetPartitioner

Partitioner that uses Dirichlet distribution to allocate samples to clients.

split_dataset(dataset: Dataset) → List[Subset]

Split a single dataset into multiple subsets, each keyed by a unique client_id.

Parameters:

dataset (Dataset) – The dataset to be split.

Returns:

A dictionary where the key is a string client_id and the

value is a Subset.

Return type:

Dict[str, Subset]

split_datasets(train_dataset: Dataset, test_dataset: Dataset) → tuple[list[torch.utils.data.dataset.Subset[Any]], list[torch.utils.data.dataset.Subset[Any]]]

Split two keyconcepts (e.g., training and testing keyconcepts) into multiple pairs of subsets, each keyed by a unique client_id.

Parameters:

• train_dataset (Dataset) – The training dataset to be split.

• test_dataset (Dataset) – The testing dataset to be split.

Returns:

A dictionary where the key is a string

client_id and the value is a tuple of two Subsets (training and testing).

Return type:

Dict[str, Tuple[Subset, Subset]]