Tensorflow Federated Tutorial | David Weissteiner

In this tutorial, we implement a basic federated learning infrastructure using tensorflow_federated in Python.

This tutorial is based on TensorFlow Federated.

First, install tensorflow_federated via pip:

pip install tensorflow-federated

You can test if tensorflow_federated has been installed successfully by executing the following command:

python -c "import tensorflow_federated as tff; print(tff.federated_computation(lambda: 'Hello World')())"

Note that this command does not necessarily work for tensorflow_federated version 0.73.0 and later.

Data Partitioning

In a real world use case, each individual federated worker owns a private dataset. Since we have one complete dataset, we need to partition the data across our federated workers in order to simulate workers with distinct data. Therefore, we implement three horizontal partitioning schemes: Range partitioning, random partitioning, and round-robin partitioning.

def partitionTrainData(train, config):
    n_workers = config["num_workers"]
    match config["part_scheme"]:
        case PartitioningScheme.RANGE:
            train_parts = partitionTrainDataRange(train, n_workers)
            return train_parts
        case PartitioningScheme.RANDOM:
            train.shuffle(train.cardinality(), seed=config["seed"])
            train_parts = partitionTrainDataRange(train, n_workers)
            return train_parts
        case PartitioningScheme.ROUND_ROBIN:
            train_parts = [train.shard(n_workers, w_idx) for w_idx in range(n_workers)]
            return train_parts

def partitionTrainDataRange(train, n_workers):
    n_rows = train.cardinality().numpy()
    distribute_remainder = lambda idx: 1 if idx < (n_rows % n_workers) else 0
    train_parts = list()
    num_elements = 0
    for w_idx in range(n_workers):
        train.skip(num_elements)
        num_elements = (n_rows // n_workers) + distribute_remainder(w_idx)
        train_parts.append(train.take(num_elements))
    return train_parts

Data Partitioning

Enjoy Reading This Article?