feat(datasets) Add DistributionPartitioner to Flower Datasets

[chongshenng]

Issue

Some FL researchers use specific NIID distributions such as a pathological power law split but Flower Datasets do not currently support such partitioning schemes.

Description

The power law splitting procedure are used in several notable research such as fedprox and tamuna. To ensure that Flower Datasets can be used with Flower baselines, this partitioner needs to be developed.

Related issues/PRs

Proposal

Add a partitioner that closely reproduce the power-law partitioning scheme.

Explanation

This partitioner will include the functionalities to:

1. Accept sample counts come from a user-specified distribution.
2. Specify minimum occurrence per label per client.
3. Specify number of labels per client.

Point 1 is intentional so that other distributions (other than a log-normal distribution) can also be flexibly prescribed to this partitioning scheme. It retains its original pathological definition, i.e. partitions only have num_unique_labels_per_partition. All samples from the dataset are exhausted during sampling, if the rescale parameter is set to True.

This implementation is inspired from Li et al. "Federated Optimization in Heterogeneous Networks" (2020) https://arxiv.org/abs/1812.06127.

Visualizations

The histogram shows the pathological partitioning of a log-normal distribution for 20 partitions with the MNIST dataset, with 5 preassigned number of samples per label, and 2 unique labels per partition:

To assess the original power law implementation by Li et al. and the Flower Datasets implementation, we plot and compare the distributions from both implementation below:

The plot above uses the original configuration of 1_000 partitions, 5 preassigned number of samples per label, and 2 unique labels per partition. Our Flower Datasets implementation closely matches the original distribution and the $y=x^{-2}$ curve, which validates our implementation.

Checklist

• Implement proposed change
• Write tests
• Make CI checks pass
• Ping maintainers on Slack (channel #contributions)

Any other comments?

[adam-narozniak]

lgtm!