xain-sdk 0.4.0

XAIN is an open source framework for federated learning.

XAIN SDK

Overview

The XAIN project is building a privacy layer for machine learning so that AI projects can meet compliance such as GDPR and CCPA. The approach relies on Federated Learning as enabling technology that allows production AI applications to be fully privacy compliant.

Federated Learning also enables different use-cases that are not strictly privacy related such as connecting data lakes, reaching higher model performance in unbalanced datasets and utilising AI models on the edge.

The main components:

• Coordinator: The entity that manages all aspects of the execution of rounds for Federated Learning. This includes the registration of clients, the selection of participants for a given round, the determination of whether sufficiently many participants have sent updated local models, the computation of an aggregated global model, and the sending of the latter model to storage or other entities.
• Participant: An entity that is the originator of a local dataset that can be selected for local training in the Federated Learning.
• Selected Participant: A Participant that has been selected by the Coordinator to participate in the next or current round.
• SDK: The library which allows Participants to interact with the Coordinator and the XAIN Platform.

The source code in this project implements the XAIN SDK to provide your local application a way to communicate with the XAIN Coordinator.

Getting started

Run the XAIN Coordinator

There are two options to run the XAIN Coordinator to perform Federated Learning on locally trained models:

• Go to the main page of the project and request a demo for the XAIN Platform.
• For the self-hosted solution, see XAIN FL Project for more details.

Integrate the XAIN SDK into your project

1. Install the XAIN SDK

To install the XAIN SDK package on your machine, simply run in your terminal:

pip install xain-sdk

2. Register your application and the device to participate in the aggregation

Now you can register your Participants to participate in the Federated Learning rounds. To do so, just send the registration request to the XAIN Coordinator:

participant.py

from typing import Dict, List, Tuple
from numpy import ndarray
from xain_sdk.participant import Participant

class MyParticipant(Participant):
    def __init__(self):

        super(MyParticipant, self).__init__()

        # define or load a model to be trained
        ...

        # define or load data to be trained on
        ...

    def train_round(
        self, weights: List[ndarray], epochs: int, epoch_base: int
    ) -> Tuple[List[ndarray], int, Dict[str, ndarray]]:

        number_samples: int
        metrics: Dict[str, ndarray]

        if weights:
            # load weights into the model
            ...

            # train the model for the specified number of epochs
            weights = ...

            # gather the number of training samples and the metrics of the trained epochs
            number_samples = ...
            metrics = ...

        else:
            # initialize new weights for the model
            weights = ...
            number_samples = 0
            metrics = {}

        # return the updated weights, number of train samples and gathered metrics
        return weights, number_samples, metrics

start.py

from xain_sdk.participant_state_machine import start_participant

# Import MyParticipant from your participant.py file 
from participant import MyParticipant

# Create a new participant
p = MyParticipant()

# Register your new participant to interact with XAIN Coordinator (hosted at XAIN Platform or self-hosted solution). The function start_participant requires two arguments:
#   - your new participant to register to interact with Coordinator,
#   - the URL of the Coordinator to connect to. 
start_participant(p, "your_host:your_port")

Now you have registered a participant. Simply repeat this step for all the participants you wish to register.

The XAIN Coordinator will take care of the rest:

• The aggregation of your locally pretrained models.
• Triggering new training rounds for the selected participants and aggregating these models.

Model metrics

A monitoring feature, which will be available as a XAIN Platform solution. If you would like to compare the performance of aggregated models, please send the specific metrics of your use case that you wish to monitor to the XAIN Coordinator. This will then be reflected in the web interface under the Project Management tab. In order to send your metrics to the XAIN Coordinator, you will need to update the train_round() method accordingly.

The returned metrics dictionary expects strs as chosen names of the metrics, which are mapped to ndarrays as values of the metrics. The first dimension of the ndarray should refer to the trained epochs of the current round and the following dimensions to the metric values per epoch. For example, a scalar-sized metric like the loss will result in a ndarray of shape (epochs, 1) and a vector-sized metric like accuracy per category will result in a ndarray of shape (epochs, number_of_categories). But in general, any metric of any positive dimension will be accepted. The metrics object then might for example look like

metrics = {
    "loss": loss_array,
    "accuracy": accuracy_array,
    "accuracy_per_category": accuracy_per_category_array,
}

Utility

The Participant base class provide some utility methods to help with the implementation of the train_round() method, namely:

• set_tensorflow_weights(): Set the weights of a Tensorflow model from a flat weight vector.
• get_tensorflow_weights(): Get and flatten the weights of a Tensorflow model.
• set_pytorch_weights(): Set the weights of a Pytorch model from a flat weight vector.
• get_pytorch_weights(): Get and flatten the weights of a Pytorch model.

Examples

Please see the following examples showing how to implement your Participant with the SDK:

• Keras/Tensorflow example for the SDK Participant implementation
• PyTorch example for the SDK Participant implementation

Testing

You can connect multiple participants at once running in parallel to a coordinator with the following script:

Bash script for starting multiple participants

Getting help

If you get stuck, don't worry, we are here to help. You can find more information here:

• More information about the project
• SDK Documentation
• GitHub issues
• More information about XAIN

In case you don't find answers to your problems or questions, simply ask your question to the community here:

• Gitter XAIN PROJECT