feat(datasets) Add DistributionPartitioner to Flower Datasets (#3791)

adap · Jul 24, 2024 · 7dc7c8a · 7dc7c8a
1 parent be114b6
commit 7dc7c8a
Show file tree

Hide file tree

Showing 5 changed files with 610 additions and 0 deletions.
diff --git a/datasets/README.md b/datasets/README.md
@@ -41,6 +41,7 @@ Create **custom partitioning schemes** or choose from the **implemented [partiti
 * Partitioner (the abstract base class) `Partitioner`
 * IID partitioning `IidPartitioner(num_partitions)`
 * Dirichlet partitioning `DirichletPartitioner(num_partitions, partition_by, alpha)`
+* Distribution partitioning `DistributionPartitioner(distribution_array, num_partitions, num_unique_labels_per_partition, partition_by, preassigned_num_samples_per_label, rescale)`
 * InnerDirichlet partitioning `InnerDirichletPartitioner(partition_sizes, partition_by, alpha)`
 * Pathological partitioning `PathologicalPartitioner(num_partitions, partition_by, num_classes_per_partition, class_assignment_mode)`
 * Natural ID partitioning `NaturalIdPartitioner(partition_by)`

diff --git a/datasets/doc/source/index.rst b/datasets/doc/source/index.rst
@@ -93,6 +93,7 @@ Here are a few of the ``Partitioner`` s that are available: (for a full list see
 * Partitioner (the abstract base class) ``Partitioner``
 * IID partitioning ``IidPartitioner(num_partitions)``
 * Dirichlet partitioning ``DirichletPartitioner(num_partitions, partition_by, alpha)``
+* Distribution partitioning ``DistributionPartitioner(distribution_array, num_partitions, num_unique_labels_per_partition, partition_by, preassigned_num_samples_per_label, rescale)``
 * InnerDirichlet partitioning ``InnerDirichletPartitioner(partition_sizes, partition_by, alpha)``
 * PathologicalPartitioner ``PathologicalPartitioner(num_partitions, partition_by, num_classes_per_partition, class_assignment_mode)``
 * Natural ID partitioner ``NaturalIdPartitioner(partition_by)``

diff --git a/datasets/flwr_datasets/partitioner/__init__.py b/datasets/flwr_datasets/partitioner/__init__.py
@@ -16,6 +16,7 @@
 
 
 from .dirichlet_partitioner import DirichletPartitioner
+from .distribution_partitioner import DistributionPartitioner
 from .exponential_partitioner import ExponentialPartitioner
 from .iid_partitioner import IidPartitioner
 from .inner_dirichlet_partitioner import InnerDirichletPartitioner
@@ -29,6 +30,7 @@
 
 __all__ = [
     "DirichletPartitioner",
+    "DistributionPartitioner",
     "ExponentialPartitioner",
     "IidPartitioner",
     "InnerDirichletPartitioner",