A Quick Start on Your Journey to Federated Learning

In my earlier post, I described the importance of federal learning from a data scientist’s view. I will now get you started on FL using your own datasets. There are several FL frameworks available, along with tutorials and user guides. However, adapting these frameworks on your own datasets is not a simple task. In this article, I will provide you the concise solution to start your FL journey with one of the popular frameworks and that is OpenFL.

Base Template

OpenFL has provided good tutorials for learning FL. We will use their Keras MNIST tutorial as our starting point. This will be our template for further experimentation. I suggest you run this tutorial on your machine to ensure the FL setup. Though lots of customization is possible, I will give you the key areas to change in this template for kick-starting your FL journey. Specifically, I will show you:

Using another image dataset instead of MNIST.
How to set up your own ANN architecture? This is required for optimizing model performance for your own datasets.
How to set up the number of collaborators (participants) for federated learning?

Using CIFAR in place of MNIST

The moment you decide to use another dataset in the template, the code for loading and processing data would change. Like MNIST, CIFAR is available as a built-in dataset in Keras. So, loading this dataset into your project is trivial. Simply import it in your project and call _loaddata method.

from keras.datasets import cifar10
(X_train, y_train), (X_test, y_test) =cifar10.load_data()

The CIFAR images have dimensions 32×32, while MNIST has dimensions 28×28. In the data preparation, you may scale the images to the sizes as required by your ANN input.

for index in range(len(images)):
img_list.append(np.resize(tf.image.rgb_to_grayscale(images[index]),(32,32)))

CIFAR has the same number of classes as MNIST. So, there is no need to change the value of the classes variable:

classes = 10

Now, as you have done with changes for loading and preprocessing the CIFAR dataset, I will show you how to change the network architecture.

Setting up Network Architecture

The MNIST tutorial sets up the network architecture in the _buildmodel method. We change it to plug-in our own architecture.

from tensorflow.keras.layers import MaxPool2D
from tensorflow.keras.layers import Dense, Dropout, Flatten, Conv2D
from tensorflow.keras.layers import MaxPool2D, Activation, MaxPooling2D

def build_model(input_shape,classes):
  model = Sequential()
  model.add(Conv2D(16, kernel_size = 3,kernel_initializer='he_normal', activation='relu', input_shape = (32, 32, 1)))
  model.add(Conv2D(32, kernel_size = 3,kernel_initializer='he_normal', activation='relu'))
  model.add(Dropout(0.2))
  model.add(Conv2D(64, kernel_size = 3, kernel_initializer='he_normal', strides=1, activation='relu'))
  model.add(MaxPooling2D((2, 2)))
  model.add(Conv2D(128, kernel_size = 3, strides=1, kernel_initializer='he_normal' ,padding='same', activation='relu'))
  model.add(MaxPooling2D((2, 2)))
  model.add(Conv2D(256, kernel_size = 3,kernel_initializer='he_normal', activation='relu'))
  model.add(MaxPooling2D((4, 4)))
  model.add(Dropout(0.2))
  model.add(Flatten())
  model.add(Dense(100,kernel_initializer='he_normal', activation = "relu"))
  model.add(Dropout(0.1))
  model.add(Dense(10, kernel_initializer='glorot_uniform', activation = "softmax"))
  # Compile the model
  model.compile(loss="categorical_crossentropy", optimizer="adam", metrics=['accuracy'])
  return model

The network plot is shown here for your quick visualization.

You will typically set your own architectures for better performance and also the type of problem that you are handling.

Finally, I will show you how to set up the number of collaborators for your model training.

Setting up Number of Collaborators

In federated learning, the most important part is to set up the number of participants who will contribute to the model training. We simply do this in a few lines of code.

We set the number of collaborators in the call to the setup method.

collaborator_models = fl_model.setup(num_collaborators=5)

Next, you need to set up the collaborators array for the above specified 5 collaborators.

collaborators = {'one':collaborator_models[0],'two':collaborator_models[1], 'three':collaborator_models[2],'four':collaborator_models[3],'five':collaborator_models[4]}

The data loader will allocate an appropriate portion of the entire dataset to each model. You can check the data allocated for model 4 using the code snippet here:

#Collaborator five's data
print(f'Collaborator five's training data size: {len(collaborator_models[4].data_loader.X_train)}')
print(f'Collaborator five's validation data size: {len(collaborator_models[4].data_loader.X_valid)}n')

You will see the following output:

Collaborator five's training data size: 9000
Collaborator five's validation data size: 1000

You are now all set for model training using federated learning.

Model Training

In the call to the training method, you can set the number of rounds to your desired value. I have set it to 20 rounds in my experiment.

final_fl_model = fx.run_experiment(collaborators,override_config={'aggregator.settings.rounds_to_train':20,"data_loader.settings.collaborator_count": 5})

Incidentally, the plan.yaml file in your template contains several parameters, which you can override in the _runexperiment method call. In the above statement, I have overwritten the defaults for _rounds_totrain and _collaboratorcount.

It takes a while to train the model. After the training halts, save the model and evaluate its performance.

The entire project source is available on my GitHub.

What’s Next?

To use your own datasets, load it into the project, do the preprocessing, set up the network and collaborators. Train the model when all is done. We did all this in a simulated environment.

How do you share the data and the plan in a real world situation?

You need to share the FL plan and model code manually with each participant. We do this using an export command in OpenFL.

How are training weights shared?

During federated training, the communication between the collaborators and the aggregator running in the OpenFL backend takes place with the help of remote procedure calls (RPC). As per the pre-defined plan, the server assigns specific tasks to a collaborator along with the dependencies and the initial data. When the collaborator finishes its training, it shares the updated weights with the aggregator through the RPC. The aggregator combines the updates received from various collaborators into a global model using the algorithm specified in the FL plan. Server then shares the new weights with the collaborators for the next round of training. The model improvement process can continue forever or until we achieve a predetermined level of accuracy.

Conclusions

Federated learning plays a vital role in Machine Learning when data privacy is of most importance. There are several FL frameworks available. You now have a quick start on your FL journey using one such popular framework – OpenFL. In this article, I have shown you how to use the base template provided by OpenFL in a dataset of your choice, how to set up a custom ANN, and how to set the number of participants. You still need to go deeper into the APIs of FL framework for a successful practical implementation. This white paper may provide further help. Good luck!

The entire project source is available on my GitHub.