Partition Data using Dirichlet Distribution

Dirichlet partitioning is a common technique to split labelled data into heterogeneous partitions with a configurable degree of class-imbalance (i.e., heterogeneity). Since a 5 minute web search did not lead me to a readily implemented solution, this guide provides tensorflow implementation of Dirichlet partitioning with a short explanation of important code chunks.

Note that this implementation is based on FedDyn. Our code is also available at GitHub.

Setup

We start by installing the required libraries via pip:

import numpy as np
import pandas as pd
import queue
import tensorflow as tf

Suppose that the following variables are the inputs to our partitioning function, thus they are already defined:

data: The data as type `tf.data.Dataset’ containing features and labels. See Tensorflow Tutorial: Load Custom Image Dataset for an example.
n_partitions: The desired number of resulting partitions.
n_classes: The number of unique classes/labels in the dataset.
alpha: The parameter for the Dirichlet distribution. Controls the degree of heterogeneity.
seed: A seed for reproducibility.

First, we set the seed and derive some helper variables:

np.random.seed(seed)
n_rows = data.cardinality().numpy()

def partitionDataDirichlet(data, n_partitions, n_classes, alpha, seed):
    np.random.seed(seed)
    iter_data = data.as_numpy_iterator()
    class_data_indices = dict()
    class_lookup_dict = dict()
    newclass_index = 0
    for idx, elem in enumerate(iter_data):
        response = elem[1]
        if(not class_lookup_dict or
            not np.any(np.all(response == np.array(list(class_lookup_dict.values())), axis=1))):
            class_lookup_dict[newclass_index] = response
            class_data_indices[newclass_index] = queue.Queue()
            newclass_index += 1
        class_data_indices[next(filter(lambda key: np.all(class_lookup_dict[key] == response),
            class_lookup_dict))].put(idx)

    n_rows = data.cardinality().numpy()
    distribute_remainder = lambda idx: 1 if idx < (n_rows % n_partitions) else 0
    # balanced cardinality
    partition_num_elements = {p_idx: (n_rows // n_partitions) + distribute_remainder(p_idx)
        for p_idx in range(n_partitions)}

    # draw the class priors for all partitions from a dirichlet distribution
    class_priors = np.random.dirichlet(alpha=[alpha]*n_classes, size=n_partitions)
    partition_assigned_classes = dict()
    partition_assigned_indices = dict()
    for _ in range(n_rows):
        # randomly select the partition to which to add an element
        current_partition = np.random.choice([p_idx for p_idx, n_elements
            in partition_num_elements.items() if n_elements > 0])
        current_priors = [class_priors[current_partition][c_idx] * (not class_data_indices[c_idx].empty())
                for c_idx in class_data_indices.keys()]
        current_priors = np.array(current_priors) / np.sum(current_priors)
        assign_class = np.random.choice(list(class_data_indices.keys()),
            p=current_priors)
        partition_assigned_classes[current_partition][assign_class] = partition_assigned_classes.setdefault(
            current_partition, dict()).setdefault(assign_class, 0) + 1
        partition_assigned_indices.setdefault(current_partition, list()).append(
            class_data_indices[assign_class].get())
    partition_assigned_indices = {p_idx: partition_assigned_indices[p_idx] for p_idx in range(n_partitions)}

    partitions = list()
    for p_idx, pai in partition_assigned_indices.items():
        part_x, part_y = zip(*list(map(list(data).__getitem__, pai)))
        part_x = np.array(part_x)
        part_y = np.array(part_y)
        part_x = tf.data.Dataset.from_tensor_slices(part_x)
        part_y = tf.data.Dataset.from_tensor_slices(part_y)
        partitions.append(tf.data.Dataset.zip((part_x, part_y)))

    print(f'Partitioned data into {n_partitions} partitions with ' +
        f'{list(partition_num_elements.values())} respectively with the following ' +
        "distribution of target labels:\n" +
        pd.DataFrame.from_dict(partition_assigned_classes).sort_index().to_string())

    return partitions

Unfinished.

Setup

Enjoy Reading This Article?