Add shard partitioner

datasets/flwr_datasets/partitioner/__init__.py

@@ -20,6 +20,7 @@
 from .linear_partitioner import LinearPartitioner
 from .natural_id_partitioner import NaturalIdPartitioner
 from .partitioner import Partitioner
+from .shard_partitioner import ShardPartitioner
 from .size_partitioner import SizePartitioner
 from .square_partitioner import SquarePartitioner
 
@@ -30,5 +31,6 @@
     "SizePartitioner",
     "LinearPartitioner",
     "SquarePartitioner",
+    "ShardPartitioner",
     "ExponentialPartitioner",
 ]

datasets/flwr_datasets/partitioner/shard_partitioner.py

@@ -0,0 +1,290 @@
+# Copyright 2023 Flower Labs GmbH. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Shard partitioner class."""
+# pylint: disable=R0912
+import math
+from typing import Dict, List, Optional
+
+import numpy as np
+
+import datasets
+from flwr_datasets.partitioner.partitioner import Partitioner
+
+
+class ShardPartitioner(Partitioner):  # pylint: disable=R0902
+    """Partitioner based on shard of (typically) unique classes.
+
+    The algorithm works as follows: the dataset is sorted by label e.g. [samples with
+    label 1, samples with labels 2 ...], then the shards are created, with each
+    shard of size - shard_size if provided or automatically calculated:
+    shards_size = len(dataset) / num_partitions * num_shards_per_node.
+    Each partition is created from `num_shards_per_node` that are chosen randomly.
+
+    There are a few ways of partitioning data that result in certain properties
+    (depending on the parameters specification):
+    1) same number of shards per nodes + the same shard size (specify:
+    a) num_shards_per_nodes, shard_size or b) num_shards_per_node)
+    In case of b the shard_size is calculated as floor(len(dataset) /
+    (num_shards_per_nodes * num_partitions))
+    2) possibly different number of shards per node (use nearly all data) + the same
+    shard size (specify: shard_size + keep_incomplete_shard=False)
+    3) possibly different number of shards per node (use all data) + possibly different
+    shard size (specify: shard_size + keep_incomplete_shard=True)
+
+
+    Algorithm based on the description in  Communication-Efficient Learning of Deep
+    Networks from Decentralized Data https://arxiv.org/abs/1602.05629. This
+    implementation expands on the initial idea by enabling more hyperparameters
+    specification therefore providing more control of the partitions. It enables the
+    division obtained in original paper.
+
+    Parameters
+    ----------
+    num_partitions : int
+        The total number of partitions that the data will be divided into.
+    partition_by : str
+        Column name of the labels (targets) based on which Dirichlet sampling works.
+    num_shards_per_node : Optional[int]
+        Number of shards to assign to a single partitioner. It's an alternative to
+        num_partitions.
+    shard_size : Optional[int]
+        Size of a single shards (a partition has one or more shards). If the size is not
+        given it will be automatically computed such that.
+    keep_incomplete_shard : bool
+        Weather to drop the last shard which might be incomplete (smaller than the
+        others). If it is dropped each shard is equal size. (It does not mean that each
+        client gets equal number of shards, which only happens if
+        num_partitions % num_shards = 0). This parameter has no effect if
+        num_shards_per_nodes and shard_size are specified.
+    shuffle: bool
+        Whether to randomize the order of samples. Shuffling applied after the
+        samples assignment to nodes.
+    seed: int
+        Seed used for dataset shuffling. It has no effect if `shuffle` is False.
+
+    Examples
+    --------
+    1) If you need same number of shards per nodes + the same shard size (and you know
+    both of these values)
+    >>> from flwr_datasets import FederatedDataset
+    >>> from flwr_datasets.partitioner import ShardPartitioner
+    >>>
+    >>> partitioner = ShardPartitioner(num_partitions=10, partition_by="label",
+    >>>                                num_shards_per_node=2, shard_size=1_000)
+    >>> fds = FederatedDataset(dataset="mnist", partitioners={"train": partitioner})
+    >>> partition = fds.load_partition(0)
+    >>> print(partition[0])  # Print the first example
+    {'image': <PIL.PngImagePlugin.PngImageFile image mode=L size=28x28 at 0x15F616C50>,
+    'label': 3}
+    >>> partition_sizes = [len(fds.load_partition(node_id)) for node_id in range(10)]
+    >>> print(partition_sizes)
+    [2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000, 2000]
+
+    2) If you want to use nearly all the data and do not need to have the number of
+    shard per each node
+    >>> from flwr_datasets import FederatedDataset
+    >>> from flwr_datasets.partitioner import ShardPartitioner
+    >>>
+    >>> partitioner = ShardPartitioner(num_partitions=9, partition_by="label",
+    >>>                                shard_size=1_000)
+    >>> fds = FederatedDataset(dataset="mnist", partitioners={"train": partitioner})
+    >>> partition_sizes = [len(fds.load_partition(node_id)) for node_id in range(9)]
+    >>> print(partition_sizes)
+    [7000, 7000, 7000, 7000, 7000, 7000, 6000, 6000, 6000]
+
+    3) If you want ot use all the data
+        >>> from flwr_datasets import FederatedDataset
+    >>> from flwr_datasets.partitioner import ShardPartitioner
+    >>>
+    >>> partitioner = ShardPartitioner(num_partitions=10, partition_by="label",
+    >>>                                shard_size=990, keep_incomplete_shard=True)
+    >>> fds = FederatedDataset(dataset="mnist", partitioners={"train": partitioner})
+    >>> partition_sizes = [len(fds.load_partition(node_id)) for node_id in range(10)]
+    >>> print(sorted(partition_sizes))
+    [5550, 5940, 5940, 5940, 5940, 5940, 5940, 5940, 5940, 6930]
+    """
+
+    def __init__(  # pylint: disable=R0913
+        self,
+        num_partitions: int,
+        partition_by: str,
+        num_shards_per_node: Optional[int] = None,
+        shard_size: Optional[int] = None,
+        keep_incomplete_shard: bool = False,
+        shuffle: bool = True,
+        seed: Optional[int] = 42,
+    ) -> None:
+        super().__init__()
+        # Attributes based on the constructor
+        self._num_partitions = num_partitions
+        self._check_num_partitions_greater_than_zero()
+        self._partition_by = partition_by
+        self._num_shards_per_node = num_shards_per_node
+        self._total_num_shards: Optional[int] = None
+        self._shard_size = shard_size
+        self._keep_incomplete_shard = keep_incomplete_shard
+        self._shuffle = shuffle
+        self._seed = seed
+
+        # Utility attributes
+        self._rng = np.random.default_rng(seed=self._seed)  # NumPy random generator
+        self._node_id_to_indices: Dict[int, List[int]] = {}
+        self._node_id_to_indices_determined = False
+
+    def load_partition(self, node_id: int) -> datasets.Dataset:
+        """Load a partition based on the partition index.
+
+        Parameters
+        ----------
+        node_id : int
+            the index that corresponds to the requested partition
+
+        Returns
+        -------
+        dataset_partition : Dataset
+            single partition of a dataset
+        """
+        # The partitioning is done lazily - only when the first partition is
+        # requested. Only the first call creates the indices assignments for all the
+        # partition indices.
+        self._check_num_partitions_correctness_if_needed()
+        self._check_possibility_of_partitions_creation()
+        self._sort_dataset_if_needed()
+        self._determine_node_id_to_indices_if_needed()
+        return self.dataset.select(self._node_id_to_indices[node_id])
+
+    def _determine_node_id_to_indices_if_needed(self) -> None:  # pylint: disable=R0914
+        """Assign sample indices to each node id.
+
+        This method works on sorted datasets. "Shard" is a part of the dataset of the
+        consecutive samples (if self._keep_incomplete_shard is False, each shard is same
+        size).
+        """
+        if self._node_id_to_indices_determined:
+            return
+        if self._num_shards_per_node is not None:
+            self._total_num_shards = int(
+                self._num_partitions * self._num_shards_per_node
+            )
+            num_shards_per_node_array = (
+                np.ones(self._num_partitions) * self._num_shards_per_node
+            )
+            self._compute_shard_size_if_missing()
+        elif self._num_shards_per_node is None:
+            if self._shard_size is None:
+                raise ValueError(
+                    "The shard_size needs to be specified if the "
+                    "num_shards_per_node is None"
+                )
+            if self._keep_incomplete_shard is False:
+                self._total_num_shards = int(
+                    math.floor(len(self.dataset) / self._shard_size)
+                )
+            elif self._keep_incomplete_shard is True:
+                self._total_num_shards = int(
+                    math.ceil(len(self.dataset) / self._shard_size)
+                )
+            else:
+                raise ValueError(
+                    "The keep_incomplete_shards need to be specified "
+                    "when _num_shards_per_node is None."
+                )
+            num_shards_per_node = int(self._total_num_shards / self._num_partitions)
+            # Assign the shards per nodes (so far, the same as in ideal case)
+            num_shards_per_node_array = (
+                np.ones(self._num_partitions) * num_shards_per_node
+            )
+            num_shards_assigned = self._num_partitions * num_shards_per_node
+            num_shards_to_assign = self._total_num_shards - num_shards_assigned
+            # Assign the "missing" shards
+            for i in range(num_shards_to_assign):
+                num_shards_per_node_array[i] += 1
+
+        else:
+            raise ValueError(
+                "The specification of nm_shards_per_node and "
+                "keep_incomplete_shards is not correct."
+            )
+
+        indices_on_which_to_split_shards = np.cumsum(
+            num_shards_per_node_array, dtype=int
+        )
+        shard_indices_array = np.random.permutation(self._total_num_shards)
+        # Randomly assign shards to node_id
+        nid_to_shard_indices = np.split(
+            shard_indices_array, indices_on_which_to_split_shards
+        )[:-1]
+        node_id_to_indices: Dict[int, List[int]] = {
+            cid: [] for cid in range(self._num_partitions)
+        }
+        # Compute node_id to sample indices based on the shard indices
+        for node_id in range(self._num_partitions):
+            for shard_idx in nid_to_shard_indices[node_id]:
+                start_id = int(shard_idx * self._shard_size)
+                end_id = min(int((shard_idx + 1) * self._shard_size), len(self.dataset))
+                node_id_to_indices[node_id].extend(list(range(start_id, end_id)))
+        if self._shuffle:
+            for indices in node_id_to_indices.values():
+                # In place shuffling
+                self._rng.shuffle(indices)
+        self._node_id_to_indices = node_id_to_indices
+        self._node_id_to_indices_determined = True
+
+    def _check_num_partitions_correctness_if_needed(self) -> None:
+        """Test num_partitions when the dataset is given (in load_partition)."""
+        if not self._node_id_to_indices_determined:
+            if self._num_partitions > self.dataset.num_rows:
+                raise ValueError(
+                    "The number of partitions needs to be smaller than the number of "
+                    "samples in the dataset."
+                )
+
+    def _check_num_partitions_greater_than_zero(self) -> None:
+        """Test num_partition left sides correctness."""
+        if not self._num_partitions > 0:
+            raise ValueError("The number of partitions needs to be greater than zero.")
+
+    def _sort_dataset_if_needed(self) -> None:
+        """Sort dataset prior to determining the partitions.
+
+        Operation only needed to be performed one time. It's required for the creation
+        of shards with the same labels.
+        """
+        if self._node_id_to_indices_determined:
+            return
+        self._dataset = self.dataset.sort(self._partition_by)
+
+    def _compute_shard_size_if_missing(self) -> None:
+        """Compute the parameters needed to perform sharding.
+
+        This method should be called after the dataset is assigned.
+        """
+        if self._shard_size is None:
+            # If shard size is not specified it needs to be computed
+            num_rows = self.dataset.num_rows
+            self._shard_size = int(num_rows / self._total_num_shards)
+
+    def _check_possibility_of_partitions_creation(self) -> None:
+        if self._shard_size is not None and self._num_shards_per_node is not None:
+            implied_min_dataset_size = (
+                self._shard_size * self._num_shards_per_node * self._num_partitions
+            )
+            if implied_min_dataset_size > len(self.dataset):
+                raise ValueError(
+                    f"Based on the given arguments the creation of the "
+                    "partitions is impossible. The implied minimum dataset"
+                    f"size is {implied_min_dataset_size} but the dataset"
+                    f"size is {len(self.dataset)}"
+                )