# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python (pytorch)
#     language: python
#     name: pytorch
# ---

# # Federated learning: using a TensorFlow model
#
# This notebook is a copy of the notebook [Federated learning basic concepts](./federated_learning_basic_concepts.ipynb). The difference is that, here, the model is built by defining a custom layer. However, apart from that, the structure is identical so the text has been removed for clearness. Please refer to the original notebook for the detailed description of the experiment. 
#
# ## The data

# +
import shfl

database = shfl.data_base.Emnist()
train_data, train_labels, test_data, test_labels = database.load_data()

print(len(train_data))
print(len(test_data))
print(type(train_data[0]))
train_data[0].shape

import matplotlib.pyplot as plt

plt.imshow(train_data[0])

iid_distribution = shfl.data_distribution.IidDataDistribution(database)
federated_data, test_data, test_label = iid_distribution.get_federated_data(num_nodes=20, percent=10)

print(type(federated_data))
print(federated_data.num_nodes())
federated_data[0].private_data
# -

# ## The model

# +
import tensorflow as tf
#If you want execute in GPU, you must uncomment this two lines.
# physical_devices = tf.config.experimental.list_physical_devices('GPU')
# tf.config.experimental.set_memory_growth(physical_devices[0], True)

class CustomDense(tf.keras.layers.Layer):
    """
    Implementation of Linear layer

    Attributes
    ----------
    units : int
        number of units for the output
    w : matrix
        Weights from the layer
    b : array
        Bias from the layer
    """

    def __init__(self, units=32, **kwargs):
        super(CustomDense, self).__init__(**kwargs)
        self._units = units
        
    def get_config(self):
        config = {'units': self._units}
        base_config = super(CustomDense, self).get_config()

        return dict(list(base_config.items()) + list(config.items()))

    def build(self, input_shape):
        """
        Method for build the params

        Parameters
        ----------
        input_shape: list
            size of inputs
        """
        self._w = self.add_weight(shape=(input_shape[-1], self._units),
                                  initializer='random_normal',
                                  trainable=True)

        self._b = self.add_weight(shape=(self._units,),
                                  initializer='random_normal',
                                  trainable=True)

    def call(self, inputs):
        """
        Apply linear layer

        Parameters
        ----------
        inputs: matrix
            Input data

        Return
        ------
        result : matrix
            the result of linear transformation of the data
        """
        return tf.nn.bias_add(tf.matmul(inputs, self._w), self._b)


def model_builder():
    inputs = tf.keras.Input(shape=(28, 28, 1))
    x = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), padding='same', activation='relu', strides=1)(inputs)
    x = tf.keras.layers.MaxPooling2D(pool_size=2, strides=2, padding='valid')(x)
    x = tf.keras.layers.Dropout(0.4)(x)
    x = tf.keras.layers.Conv2D(32, kernel_size=(3, 3), padding='same', activation='relu', strides=1)(x)
    x = tf.keras.layers.MaxPooling2D(pool_size=2, strides=2, padding='valid')(x)
    x = tf.keras.layers.Flatten()(x)
    x = CustomDense(128)(x)
    x = tf.nn.relu(x)
    x = tf.keras.layers.Dropout(0.1)(x)
    x = CustomDense(64)(x)
    x = tf.nn.relu(x)
    x = CustomDense(10)(x)
    outputs = tf.nn.softmax(x)
    
    model = tf.keras.Model(inputs=inputs, outputs=outputs)

    criterion = tf.keras.losses.CategoricalCrossentropy()
    optimizer = tf.keras.optimizers.RMSprop()
    metrics = [tf.keras.metrics.categorical_accuracy]
    
    return shfl.model.DeepLearningModel(model=model, criterion=criterion, optimizer=optimizer, metrics=metrics)


# -

aggregator = shfl.federated_aggregator.FedAvgAggregator()
federated_government = shfl.federated_government.FederatedGovernment(model_builder, federated_data, aggregator)

# +
import numpy as np

class Reshape(shfl.private.FederatedTransformation):
    
    def apply(self, labeled_data):
        labeled_data.data = np.reshape(labeled_data.data, (labeled_data.data.shape[0], labeled_data.data.shape[1], labeled_data.data.shape[2],1))

class CastFloat(shfl.private.FederatedTransformation):
    
    def apply(self, labeled_data):
        labeled_data.data = labeled_data.data.astype(np.float32)
        
shfl.private.federated_operation.apply_federated_transformation(federated_data, Reshape())
shfl.private.federated_operation.apply_federated_transformation(federated_data, CastFloat())
# -

# ## Run the federated learning experiment

test_data = np.reshape(test_data, (test_data.shape[0], test_data.shape[1], test_data.shape[2],1))
test_data = test_data.astype(np.float32)
federated_government.run_rounds(3, test_data, test_label)

StarcoderdataJupyterScriptsDedupFiltered