from going_modular import *

import warnings
warnings.simplefilter("ignore")

"""
Script to train a model on a dataset using PyTorch with the classic ML approach (no federated learning).
"""
main_parser = parsing(description='PyTorch baseline Training')
args = main_parser.parse_args()
# To define the device allocation in PyTorch:
DEVICE = torch.device(choice_device(args.device))  # Try "cuda" to train on GPU
print(f"Training on {DEVICE} using PyTorch {torch.__version__}")

if __name__ == '__main__':
    # 1) training parameters`
    number_clients = args.number_clients
    epochs = args.max_epochs
    matrix_path = args.matrix_path
    CLASSES = classes_string(args.dataset)
    NUM_CLASSES = len(CLASSES)
    yaml_path = args.yaml_path

    # 2) Load model and data
    trainloaders, valloaders, testloader = data_setup.load_datasets(num_clients=number_clients,
                                                                    batch_size=args.batch_size,
                                                                    resize=args.length, seed=args.seed,
                                                                    num_workers=args.num_workers, splitter=args.split,
                                                                    dataset=args.dataset, data_path=args.data_path,
                                                                    data_path_val=args.data_path_val)

    # To get the first trainloader and valloader (for the first client)
    trainloader = trainloaders[0]
    valloader = valloaders[0]

    # to define the model and the optimizer
    net = Net(num_classes=len(CLASSES)).to(DEVICE)
    optimizer = torch.optim.SGD(net.parameters(), lr=args.lr, momentum=0.9)

    # to get the trained model and the trained parameters (optimizer, metrics, ...)
    if os.path.exists(args.model_save):
        print(" To get the checkpoint")
        checkpoint = torch.load(args.model_save, map_location=DEVICE)

        net.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        epoch = checkpoint['epoch']
        loss = checkpoint['loss']
        acc = checkpoint['acc']

    else:
        # To give bad values
        acc = 0
        loss = 1000
        epoch = 0

    # to define the loss function
    criterion = torch.nn.CrossEntropyLoss()

    # 3) Train the model
    start_train = time.time()
    results = engine.train(net, trainloader, valloader, optimizer=optimizer, loss_fn=criterion, epochs=epochs,
                           device=DEVICE)

    end_train = time.time()
    # test net with another dataset
    test_loss, test_acc, y_pred, y_true, y_proba = engine.test(net, testloader, loss_fn=criterion, device=DEVICE)
    end_test = time.time()
    print(f"Final test set performance:\n\tloss {test_loss}\n\taccuracy {test_acc} %")
    print(f"Time performance : Train = {end_train - start_train}seconds  \n\tTest = {end_test - end_train} seconds")
    write_yaml({"Train_time": end_train - start_train, "Test_time": end_test - end_train},
               yaml_path,
               data1=read_yaml(yaml_path) if os.path.exists(yaml_path) else None
               )

    # Save results
    if args.save_results:
        save_graphs(args.save_results, epochs, results)

        # Build confusion matrix
        if args.matrix_path:
            save_matrix(y_true, y_pred, args.save_results + args.matrix_path, CLASSES)

        # Build roc curve
        if args.roc_path:
            save_roc(y_true, y_proba, args.save_results + args.roc_path, NUM_CLASSES)

    if args.model_save and test_acc > acc:
        print(f"save checkpoint:\nepoch: {epochs + epoch}, \t'loss: {test_loss},\tacc: {test_acc}")

        torch.save({
            'epoch': epochs + epoch,
            'model_state_dict': net.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'loss': test_loss,
            'acc': test_acc,
        }, args.model_save)

classic.py
lcoEntreprise/flower-homomorphic_encryption
import torch.nn.functional
import flwr as fl
from going_modular import *

"""
Script to define the client side of the federated learning pipeline with Flower.
"""


# #############################################################################
# 2. Federation of the pipeline with Flower
# #############################################################################
class FlowerClient(fl.client.NumPyClient):
    """
    The client class for the federated learning pipeline with Flower.
    This class is used to define the client-side logic for federated learning with Flower.
    It inherits from fl.client.NumPyClient. The only real difference between Client and NumPyClient is
    that NumPyClient takes care of serialization and deserialization for you.

    args:
        cid: client id (int)
        net: model (torch.nn.Module)
        trainloader: trainloader (torch.utils.data.DataLoader)
        valloader: valloader (torch.utils.data.DataLoader)
        device: device (torch.device)
        batch_size: batch size (int)
        save_results: path to save the results (str)
        matrix_path: path to save the confusion matrix (str)
        roc_path: path to save the roc curve (str)
        yaml_path: path to save the yaml file (str)
        he: boolean to use the homomorphic encryption (bool)
        classes: list of classes (list)
        context_client: context client (tenseal.context.Context)
    """

    def __init__(self, cid, net, trainloader, valloader, device, batch_size, save_results, matrix_path, roc_path,
                 yaml_path, he, classes, context_client):

        # Initialize the client
        self.net = net
        self.trainloader = trainloader
        self.valloader = valloader
        self.cid = cid
        self.device = device
        self.batch_size = batch_size
        self.save_results = save_results
        self.matrix_path = matrix_path
        self.roc_path = roc_path
        self.yaml_path = yaml_path
        self.he = he
        self.classes = classes
        self.context_client = context_client

    def get_parameters(self, config):
        """
        get the model parameters and return them as a list of NumPy ndarray’s
        (which is what flwr.client.NumPyClient expects)

        args:
            config: config (dict)

        return:
            parameters: parameters (list)
        """
        print(f"[Client {self.cid}] get_parameters")
        return get_parameters2(self.net, self.context_client)

    def fit(self, parameters, config):
        """
        - Update the parameters of the local model with the parameters received from the server
        - Train the updated model on the local train dataset (x_train/y_train)
        - Return the updated model parameters and the number of examples used for training to the server (for 1 given round)

        args:
            parameters: parameters (list)
            config: config (dict)

        return:
            parameters: parameters (list)
        """
        # Read values from config
        server_round = config['server_round']
        local_epochs = config['local_epochs']
        lr = float(config["learning_rate"])

        # Use values provided by the config
        print(f'[Client {self.cid}, round {server_round}] fit, config: {config}')

        # Update local model parameters
        set_parameters(self.net, parameters, self.context_client)

        criterion = torch.nn.CrossEntropyLoss()
        optimizer = torch.optim.SGD(self.net.parameters(), lr=lr,
                                    momentum=0.9)

        # Start training
        results = engine.train(self.net, self.trainloader, self.valloader, optimizer=optimizer, loss_fn=criterion,
                               epochs=local_epochs, device=self.device)

        # Save results
        if self.save_results:
            save_graphs(self.save_results, local_epochs, results, f"_Client {self.cid}")

        return get_parameters2(self.net, self.context_client), len(self.trainloader), {}

    def evaluate(self, parameters, config):
        """
        - Update the parameters of the local model with the parameters received from the server

        - Evaluate the updated model on the local test dataset (x_test/y_test)

        - Return the local loss and accuracy to the server

        args:
            parameters: parameters (list)
            config: config (dict)

        return:
            loss: loss (float)
            num_examples: number of examples (int)
            metrics: metrics (dict)
        """
        print(f"[Client {self.cid}] evaluate, config: {config}")
        set_parameters(self.net, parameters, self.context_client)

        # Evaluate global model parameters on the local test data
        loss, accuracy, y_pred, y_true, y_proba = engine.test(self.net, self.valloader,
                                                              loss_fn=torch.nn.CrossEntropyLoss(), device=self.device)

        if self.save_results:
            os.makedirs(self.save_results, exist_ok=True)
            if self.matrix_path:
                save_matrix(y_true, y_pred, self.save_results + self.matrix_path, self.classes)

            if self.roc_path:
                save_roc(y_true, y_proba, self.save_results + self.roc_path, len(self.classes))
        # Return results, including the custom accuracy metric
        return float(loss), len(self.valloader), {"accuracy": float(accuracy)}


# The client-side execution (training, evaluation) from the server-side
def client_common(cid: str,
                  model_save: str, path_yaml: str, path_roc: str, results_save: str, path_matrix: str,
                  batch_size: str, trainloaders, valloaders, DEVICE, CLASSES,
                  he=False, secret_path="", server_path=""):
    """
    args:
        cid: client id (str)
        model_save: path to save the model (str)
        path_yaml: path to save the yaml file (str)
        path_roc: path to save the roc curve (str)
        results_save: path to save the results (str)
        path_matrix: path to save the confusion matrix (str)
        batch_size: batch size (int)
        trainloaders: trainloaders (list)
        valloaders: valloaders (list)
        DEVICE: device (torch.device)
        CLASSES: classes (list)
        he: boolean to use the homomorphic encryption (bool)
        secret_path: path to save the secret key (str)
        server_path: path to save the server public key (str)

    return:
        FlowerClient: client (FlowerClient object)

    """
    # Load data
    # Note: each client gets a different trainloader/valloader, so each client will train and evaluate
    # on their own unique dataset.
    trainloader = trainloaders[int(cid)]
    valloader = valloaders[int(cid)]

    context_client = None

    # Load model
    net = Net(num_classes=len(CLASSES)).to(DEVICE)

    # Homomorphic encryption
    if he:
        print("Run with homomorphic encryption")
        if os.path.exists(secret_path):
            # To get the existing public/private keys combination
            with open(secret_path, 'rb') as f:
                query = pickle.load(f)

            context_client = ts.context_from(query["contexte"])

        else:
            # To create the public/private keys combination
            context_client = security.context()
            with open(secret_path, 'wb') as f:  # 'ab' to add existing file
                encode = pickle.dumps({"contexte": context_client.serialize(save_secret_key=True)})
                f.write(encode)

        secret_key = context_client.secret_key()

    else:
        print("Run WITHOUT homomorphic encryption")

    # C) Update the local model with the parameters received from the server
    # to get the trained model and the trained parameters (optimizer, metrics, ...)
    if os.path.exists(model_save):
        print(" To get the checkpoint")
        checkpoint = torch.load(model_save, map_location=DEVICE)['model_state_dict']
        if he:
            print("to decrypt model")
            # To decrypt the parameters with the private key
            server_query, server_context = security.read_query(server_path)
            server_context = ts.context_from(server_context)
            for name in checkpoint:
                print(name)
                # To decrypt the parameters with the private key
                checkpoint[name] = torch.tensor(
                    security.deserialized_layer(name, server_query[name], server_context).decrypt(secret_key)
                )

        # Update network with the aggregated results
        net.load_state_dict(checkpoint)

    # Create a  single Flower client representing a single organization
    return FlowerClient(cid, net, trainloader, valloader, device=DEVICE, batch_size=batch_size,
                        matrix_path=path_matrix, roc_path=path_roc, save_results=results_save, yaml_path=path_yaml,
                        he=he, context_client=context_client, classes=CLASSES)

client.py
lcoEntreprise/flower-homomorphic_encryption
import os
import tenseal as ts

from going_modular import security

"""
Script to create the public/private keys combination.
"""


def combo_keys(client_path="secret.pkl", server_path="server_key.pkl"):
    """
    To create the public/private keys combination
    args:
        client_path: path to save the secret key (str)
        server_path: path to save the server public key (str)
    """
    context_client = security.context()
    security.write_query(client_path, {"contexte": context_client.serialize(save_secret_key=True)})
    security.write_query(server_path, {"contexte": context_client.serialize()})

    _, context_client = security.read_query(client_path)
    _, context_server = security.read_query(server_path)

    context_client = ts.context_from(context_client)
    context_server = ts.context_from(context_server)
    print("Is the client context private?", ("Yes" if context_client.is_private() else "No"))
    print("Is the server context private?", ("Yes" if context_server.is_private() else "No"))


if __name__ == '__main__':
    secret_path = "secret.pkl"
    public_path = "server_key.pkl"
    if os.path.exists(secret_path):
        # To get the existing public/private keys combination
        print("it exists")
        _, context_client = security.read_query(secret_path)

    else:
        # To create the public/private keys combination
        combo_keys(client_path=secret_path, server_path=public_path)

create_keys.py
lcoEntreprise/flower-homomorphic_encryption
import torchvision
import flwr as fl
import warnings
from going_modular import *
"""
The script with the common functions for the federated learning pipeline (client and server)
"""

warnings.simplefilter("ignore")

print("flwr", fl.__version__)
print("numpy", np.__version__)
print("torch", torch.__version__)
print("torchvision", torchvision.__version__)

# #############################################################################
# 1. Regular PyTorch pipeline: nn.Module, train, test, and DataLoader
# #############################################################################
# 1) training parameters`
main_parser2 = parsing(description='Federated Learning asset')
args = main_parser2.parse_args()
CLASSES = classes_string(args.dataset)

# In addition, we define the device allocation in PyTorch with:
DEVICE = torch.device(choice_device(args.device))
print(f"Training on {DEVICE}")

# 2) Load model and data
trainloaders, valloaders, testloader = data_setup.load_datasets(num_clients=args.number_clients,
                                                                batch_size=args.batch_size, resize=args.length,
                                                                seed=args.seed, num_workers=args.num_workers,
                                                                splitter=args.split, dataset=args.dataset,
                                                                data_path=args.data_path,
                                                                data_path_val=args.data_path_val)

federated.py
lcoEntreprise/flower-homomorphic_encryption
from .common import *
from .security import *
from .model_builder import *
from .data_setup import *
from .engine import *

going_modular/__init__.py
lcoEntreprise/flower-homomorphic_encryption
import argparse
import random
import torch
import shutil
import yaml
import matplotlib.pyplot as plt
from sklearn.metrics import roc_curve, auc, confusion_matrix
import seaborn as sn
import pandas as pd
import torch.nn.functional
from collections import OrderedDict
from .security import *


def write_yaml(data, file_write='toyaml.yml', data1=None):
    """
    A function to write YAML file

    :param data: data to write in the YAML file
    :param file_write: path to save the YAML file
    :param data1: data to add in the YAML file (if we want to add data in the YAML file without overwriting it)
    :return: data (the data to write in the YAML file)
    """

    def accumul_time(time_key, data, data1):
        if time_key in data1 and time_key in data:
            data[time_key] += data1[time_key]

        return data

    path_yaml = os.path.join(*file_write.split('/')[:-1])
    print("save yaml in ", path_yaml)
    os.makedirs(path_yaml, exist_ok=True)
    with open(file_write, 'w') as f:
        if data1:
            data = accumul_time("Train_time", data, data1)
            data = accumul_time("Test_time", data, data1)
            data = {**data1, **data}

        yaml.dump(data, f)

    return data


def read_yaml(yaml_file='config.yml'):
    """
    A function to read YAML file

    :param yaml_file: path to the YAML file
    :return: config (the data in the YAML file)
    """

    with open(yaml_file) as f:
        config = yaml.safe_load(f)

    return config


def choice_device(device):
    """
    A function to choose the device

    :param device: the device to choose (cpu, gpu or mps)
    """
    if torch.cuda.is_available() and device != "cpu":
        # on Windows, "cuda:0" if torch.cuda.is_available()
        device = "cuda:0"

    elif torch.backends.mps.is_available() and torch.backends.mps.is_built() and device != "cpu":
        """
        on Mac : 
        - torch.backends.mps.is_available() ensures that the current MacOS version is at least 12.3+
        - torch.backends.mps.is_built() ensures that the current current PyTorch installation was built with MPS activated.
        """
        device = "mps"

    else:
        device = "cpu"

    return device


def classes_string(name_dataset):
    """
    A function to get the classes of the dataset

    :param name_dataset: the name of the dataset
    :return: classes (the classes of the dataset) in a tuple
    """
    if name_dataset == "cifar":
        classes = ('plane', 'car', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

    elif name_dataset == "animaux":
        classes = ('cat', 'dog')

    elif name_dataset == "breast":
        classes = ('0', '1')

    elif name_dataset == "histo":
        classes = ('0', '1')

    else:
        print("Warning problem : unspecified dataset")
        return ()

    return classes


def parsing(description='PyTorch ImageNet Training'):
    """
    A function to work with command line arguments (argparse)

    :param description: the description of the model
    :return: argparser object (the arguments of the model)
    """
    # To define the argparse arguments
    # Create the top-level parser
    parent_parser = argparse.ArgumentParser(description="common", add_help=False)
    parent_parser.add_argument('--max_epochs', type=int, default=1)
    parent_parser.add_argument('--number_clients', type=int, default=2)
    parent_parser.add_argument('--length', type=int, default=None, help='size at the entrance of the model')
    parent_parser.add_argument('--batch_size', type=int, default=64)
    parent_parser.add_argument('--device', default='cpu', type=str,
                               help="- Choice of the device between cpu and gpu "
                                    "(cuda if compatible Nvidia and mps if on mac\n"
                                    "- The choice the output may be cpu even if you choose the gpu if the latter isn't "
                                    "compatible")
    parent_parser.add_argument('--dataset', default='cifar', help="choice of the dataset (cifar10 by default)")
    parent_parser.add_argument('--data_path', type=str, default='./data/', help='Path to the training data')
    parent_parser.add_argument('--data_path_val', type=str, default=None, help='Path to the validation dataset')
    parent_parser.add_argument('--model_save', type=str, default='', help='Path to save the central model')
    parent_parser.add_argument('--yaml_path', type=str, default='./results/results.yml',
                               help='Path to save the metrics results')
    parent_parser.add_argument('--seed', type=int, default=42)
    parent_parser.add_argument('--num_workers', type=int, default=0)
    parent_parser.add_argument('--split', default=10, type=int,
                               help='ratio (in percent) of the training dataset that will be used for the test '
                                    '(default : 10)')
    parent_parser.add_argument('--lr', default=0.001, type=float,
                               help='learning rate for the central model'
                                    '(default : 0.001)')
    parent_parser.add_argument('--he', default=False, required=False, action='store_true', dest="he",
                               help='True if we want to use the homomorphic encryption (by default : False)')
    parent_parser.add_argument('--path_keys', type=str, default="secret.pkl",
                               help='Path to get the combo private/public keys')
    parent_parser.add_argument('--path_public_key', type=str, default="server_key.pkl",
                               help='Path to get the the public key')
    parent_parser.add_argument('--path_crypted', type=str, default="server.pkl",
                               help='Path to save the crypted (and not crypted) weights')

    # Create the specific commands for the "classic ML"
    parser_ml = argparse.ArgumentParser(description="classic", add_help=False)
    parser_ml.add_argument('--matrix_path', type=str, default=None,
                           help='Path to save the confusion matrix')
    parser_ml.add_argument('--roc_path', type=str, default=None,
                           help='Path to save the roc figures')
    parser_ml.add_argument('--save_results', type=str, default=None,
                           help='Path to save the results')

    # Create the specific commands for the "server"
    parser_server = argparse.ArgumentParser(description="server", add_help=False)
    parser_server.add_argument('--frac_fit', type=float, default=1.0)
    parser_server.add_argument('--frac_eval', type=float, default=0.5)
    parser_server.add_argument('--min_fit_clients', type=int, default=2)
    parser_server.add_argument('--min_eval_clients', type=int, default=None)
    parser_server.add_argument('--min_avail_clients', type=int, default=2)

    parser_server.add_argument('--rounds', default=3, type=int,
                               help='number of rounds (default : 3)')

    # Create the specific commands for the "client"
    parser_client = argparse.ArgumentParser(description="client", add_help=False)
    parser_client.add_argument('--id_client', type=str, default=None,
                               help='client id (by default None)')

    # Create the final parser
    main_parser = argparse.ArgumentParser(description=description)

    # Create the subparser of the final parser
    service_subparsers = main_parser.add_subparsers(title="service",
                                                    dest="service_command")

    # Add specific command for each choice (classic, client, server and simulation)
    classic_subparser = service_subparsers.add_parser("run", help="classic ML",
                                                      parents=[parent_parser, parser_ml])

    # python client.py client
    client_subparser = service_subparsers.add_parser("client", help="client",
                                                     parents=[parent_parser, parser_ml, parser_client])

    # python server.py server
    server_subparser = service_subparsers.add_parser("server", help="server",
                                                     parents=[parent_parser, parser_server])

    simul_subparser = service_subparsers.add_parser("simulation", help="client",
                                                    parents=[parent_parser, parser_server, parser_ml, parser_client])

    return main_parser


# functions for the dataset creation
def supp_ds_store(path):
    """
    Delete the hidden file ".DS_Store" created on macOS

    :param path: path to the folder where the hidden file ".DS_Store" is
    """
    for i in os.listdir(path):
        if i == ".DS_Store":
            print("Deleting of the hidden file '.DS_Store'")
            os.remove(path + "/" + i)


def create_files_train_test(path_init, path_final, splitter):
    """
    Split the dataset from path_init into two datasets : train and test in path_final
    with the splitter ratio (in %). Example : if splitter = 10, 10% of the initial dataset will be in the test dataset.

    :param path_init: path of the initial dataset
    :param path_final: path of the final dataset
    :param splitter: ratio (in %) of the initial dataset that will be in the test dataset.
    """
    # Move a file from rep1 to rep2
    for classe in os.listdir(path_init):
        list_init = os.listdir(path_init + "/" + classe)
        size_test = int(len(list_init) * splitter/100)
        print("Before : ", len(list_init))
        for _ in range(size_test):
            e = random.choice(list_init)  # random choice of the path of an image
            list_init.remove(e)
            shutil.move(path_init + classe + "/" + e, path_final + classe + "/" + e)

        print("After", path_init + classe, ":", len(os.listdir(path_init + classe)))
        print(path_final + classe, ":", len(os.listdir(path_final + classe)))


def save_matrix(y_true, y_pred, path, classes):
    """
    Save the confusion matrix in the path given in argument.

    :param y_true: true labels (real labels)
    :param y_pred: predicted labels (labels predicted by the model)
    :param path: path to save the confusion matrix
    :param classes: list of the classes
    """
    # To get the confusion matrix
    cf_matrix = confusion_matrix(y_true, y_pred)

    # To normalize the confusion matrix
    cf_matrix_normalized = cf_matrix / np.sum(cf_matrix) * 10

    # To round up the values in the matrix
    cf_matrix_round = np.round(cf_matrix_normalized, 2)

    # To plot the matrix
    df_cm = pd.DataFrame(cf_matrix_round, index=[i for i in classes], columns=[i for i in classes])
    plt.figure(figsize=(12, 7))
    sn.heatmap(df_cm, annot=True)
    plt.xlabel("Predicted label", fontsize=13)
    plt.ylabel("True label", fontsize=13)
    plt.title("Confusion Matrix", fontsize=15)

    plt.savefig(path)
    plt.close()


def save_roc(targets, y_proba, path, nbr_classes):
    """
    Save the roc curve in the path given in argument.

    :param targets: true labels (real labels)
    :param y_proba: predicted labels (labels predicted by the model)
    :param path: path to save the roc curve
    :param nbr_classes: number of classes
    """
    y_true = np.zeros(shape=(len(targets), nbr_classes))  # array-like of shape (n_samples, n_classes)
    for i in range(len(targets)):
        y_true[i, targets[i]] = 1

    # Compute ROC curve and ROC area for each class
    fpr = dict()
    tpr = dict()
    roc_auc = dict()
    for i in range(nbr_classes):
        fpr[i], tpr[i], _ = roc_curve(y_true[:, i], y_proba[:, i])
        roc_auc[i] = auc(fpr[i], tpr[i])

    # Compute micro-average ROC curve and ROC area
    fpr["micro"], tpr["micro"], _ = roc_curve(y_true.ravel(), y_proba.ravel())
    roc_auc["micro"] = auc(fpr["micro"], tpr["micro"])
    # First aggregate all false positive rates
    all_fpr = np.unique(np.concatenate([fpr[i] for i in range(nbr_classes)]))

    # Then interpolate all ROC curves at this points
    mean_tpr = np.zeros_like(all_fpr)
    for i in range(nbr_classes):
        mean_tpr += np.interp(all_fpr, fpr[i], tpr[i])

    # Finally average it and compute AUC
    mean_tpr /= nbr_classes

    fpr["macro"] = all_fpr
    tpr["macro"] = mean_tpr
    roc_auc["macro"] = auc(fpr["macro"], tpr["macro"])

    # Plot all ROC curves
    plt.figure()
    plt.plot(
        fpr["micro"],
        tpr["micro"],
        label="micro-average ROC curve (area = {0:0.2f})".format(roc_auc["micro"]),
        color="deeppink",
        linestyle=":",
        linewidth=4,
    )

    plt.plot(
        fpr["macro"],
        tpr["macro"],
        label="macro-average ROC curve (area = {0:0.2f})".format(roc_auc["macro"]),
        color="navy",
        linestyle=":",
        linewidth=4,
    )

    lw = 2
    for i in range(nbr_classes):
        plt.plot(
            fpr[i],
            tpr[i],
            lw=lw,
            label="ROC curve of class {0} (area = {1:0.2f})".format(i, roc_auc[i]),
        )

    plt.plot([0, 1], [0, 1], "k--", lw=lw, label='Worst case')
    plt.xlim([0.0, 1.0])
    plt.ylim([0.0, 1.05])
    plt.xlabel("False Positive Rate")
    plt.ylabel("True Positive Rate")
    plt.title("Receiver operating characteristic (ROC) Curve OvR")  # One vs Rest
    plt.legend(loc="lower right")  # loc="best"

    plt.savefig(path)
    plt.close()


def save_graphs(path_save, local_epoch, results, end_file=""):
    """
    Save the graphs in the path given in argument.

    :param path_save: path to save the graphs
    :param local_epoch: number of epochs
    :param results: results of the model (accuracy and loss)
    :param end_file: end of the name of the file
    """
    os.makedirs(path_save, exist_ok=True)  # to create folders results
    print("save graph in ", path_save)
    # plot training curves (train and validation)
    plot_graph(
        [[*range(local_epoch)]] * 2,
        [results["train_acc"], results["val_acc"]],
        "Epochs", "Accuracy (%)",
        curve_labels=["Training accuracy", "Validation accuracy"],
        title="Accuracy curves",
        path=path_save + "Accuracy_curves" + end_file)

    plot_graph(
        [[*range(local_epoch)]] * 2,
        [results["train_loss"], results["val_loss"]],
        "Epochs", "Loss",
        curve_labels=["Training loss", "Validation loss"], title="Loss curves",
        path=path_save + "Loss_curves" + end_file)


def plot_graph(list_xplot, list_yplot, x_label, y_label, curve_labels, title, path=None):
    """
    Plot the graph of the list of points (list_xplot, list_yplot)
    :param list_xplot: list of list of points to plot (one line per curve)
    :param list_yplot: list of list of points to plot (one line per curve)
    :param x_label: label of the x axis
    :param y_label: label of the y axis
    :param curve_labels: list of labels of the curves (curve names)
    :param title: title of the graph
    :param path: path to save the graph
    """
    lw = 2

    plt.figure()
    for i in range(len(curve_labels)):
        plt.plot(list_xplot[i], list_yplot[i], lw=lw, label=curve_labels[i])

    plt.xlabel(x_label)
    plt.ylabel(y_label)
    plt.title(title)

    if curve_labels:
        plt.legend(loc="lower right")

    if path:
        plt.savefig(path)


def get_parameters2(net, context_client=None) -> List[np.ndarray]:
    """
    Get the parameters of the network
    :param net: network to get the parameters (weights and biases)
    :param context_client: context of the crypted weights (if None, return the clear weights)
    :return: list of parameters (weights and biases) of the network
    """
    if context_client:
        # Crypte of the model trained at the client for a given round (after each round the model is aggregated between
        # clients)
        encrypted_tensor = crypte(net.state_dict(), context_client)  # list of encrypted layers (weights and biases)

        return [layer.get_weight() for layer in encrypted_tensor]

    return [val.cpu().numpy() for _, val in net.state_dict().items()]


def set_parameters(net, parameters: List[np.ndarray], context_client=None):
    """
    Update the parameters of the network with the given parameters (weights and biases)
    :param net: network to set the parameters (weights and biases)
    :param parameters: list of parameters (weights and biases) to set
    :param context_client: context of the crypted weights (if None, set the clear weights)
    """
    params_dict = zip(net.state_dict().keys(), parameters)
    if context_client:
        secret_key = context_client.secret_key()
        dico = {k: deserialized_layer(k, v, context_client) for k, v in params_dict}

        state_dict = OrderedDict(
            {k: torch.Tensor(v.decrypt(secret_key)) for k, v in dico.items()}
        )

    else:
        dico = {k: torch.Tensor(v) for k, v in params_dict}
        state_dict = OrderedDict(dico)

    net.load_state_dict(state_dict, strict=True)
    print("Updated model")

going_modular/common.py
lcoEntreprise/flower-homomorphic_encryption
"""
Contains functionality for creating PyTorch DataLoaders for image classification data.
"""
from torchvision import datasets, transforms
from torch.utils.data import DataLoader, random_split

from .common import *

NUM_WORKERS = os.cpu_count()

# Normalization values for the different datasets
NORMALIZE_DICT = {
    'mnist': dict(mean=(0.1307,), std=(0.3081,)),
    'cifar': dict(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)),
    'animaux': dict(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    'breast': dict(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    'histo': dict(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
    }


def split_data_client(dataset, num_clients, seed):
    """
    This function is used to split the dataset into train and test for each client.
    :param dataset: the dataset to split (type: torch.utils.data.Dataset)
    :param num_clients: the number of clients
    :param seed: the seed for the random split
    """
    # Split training set into `num_clients` partitions to simulate different local datasets
    partition_size = len(dataset) // num_clients
    lengths = [partition_size] * (num_clients - 1)
    lengths += [len(dataset) - sum(lengths)]
    ds = random_split(dataset, lengths, torch.Generator().manual_seed(seed))
    return ds


# Define model, architecture and dataset
# The DataLoaders downloads the training and test data that are then normalized.
def load_datasets(num_clients: int, batch_size: int, resize: int, seed: int, num_workers: int, splitter=10,
                  dataset="cifar", data_path="./data/", data_path_val=""):
    """
    This function is used to load the dataset and split it into train and test for each client.
    :param num_clients: the number of clients
    :param batch_size: the batch size
    :param resize: the size of the image after resizing (if None, no resizing)
    :param seed: the seed for the random split
    :param num_workers: the number of workers
    :param splitter: percentage of the training data to use for validation. Example: 10 means 10% of the training data
    :param dataset: the name of the dataset in the data folder
    :param data_path: the path of the data folder
    :param data_path_val: the absolute path of the validation data (if None, no validation data)
    :return: the train and test data loaders
    """

    list_transforms = [transforms.ToTensor(), transforms.Normalize(**NORMALIZE_DICT[dataset])]

    if dataset == "cifar":
        # Download and transform CIFAR-10 (train and test)
        transformer = transforms.Compose(
            list_transforms
        )
        trainset = datasets.CIFAR10(data_path + dataset, train=True, download=True, transform=transformer)
        testset = datasets.CIFAR10(data_path + dataset, train=False, download=True, transform=transformer)

    else:
        if resize is not None:
            list_transforms = [transforms.Resize((resize, resize))] + list_transforms

        transformer = transforms.Compose(list_transforms)
        supp_ds_store(data_path + dataset)
        supp_ds_store(data_path + dataset + "/train")
        supp_ds_store(data_path + dataset + "/test")
        trainset = datasets.ImageFolder(data_path + dataset + "/train", transform=transformer)
        testset = datasets.ImageFolder(data_path + dataset + "/test", transform=transformer)

    print(f"The training set is created for the classes : {trainset.classes}")

    # Split training set into `num_clients` partitions to simulate different local datasets
    datasets_train = split_data_client(trainset, num_clients, seed)
    if data_path_val:
        valset = datasets.ImageFolder(data_path_val, transform=transformer)
        datasets_val = split_data_client(valset, num_clients, seed)

    # Split each partition into train/val and create DataLoader
    trainloaders = []
    valloaders = []
    for i in range(num_clients):
        if data_path_val:
            # if we already have a validation dataset
            trainloaders.append(DataLoader(datasets_train[i], batch_size=batch_size, shuffle=True))
            valloaders.append(DataLoader(datasets_val[i], batch_size=batch_size))

        else:
            len_val = int(len(datasets_train[i]) * splitter / 100)  # splitter % validation set
            len_train = len(datasets_train[i]) - len_val
            lengths = [len_train, len_val]
            ds_train, ds_val = random_split(datasets_train[i], lengths, torch.Generator().manual_seed(seed))
            trainloaders.append(DataLoader(ds_train, batch_size=batch_size, shuffle=True))
            valloaders.append(DataLoader(ds_val, batch_size=batch_size))

    testloader = DataLoader(testset, batch_size=batch_size)
    return trainloaders, valloaders, testloader

going_modular/data_setup.py
lcoEntreprise/flower-homomorphic_encryption
"""
Contains functions for training and testing a PyTorch model.
"""
import torch
import torch.nn as nn
from tqdm.auto import tqdm
from typing import Dict, List, Tuple
import numpy as np


def test(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, loss_fn: torch.nn.Module,
         device: torch.device):
    """Tests a PyTorch model for a single epoch.

    Turns a target PyTorch model to "eval" mode and then performs
    a forward pass on a testing dataset.

    Args:
    model: A PyTorch model to be tested.
    dataloader: A DataLoader instance for the model to be tested on.
    loss_fn: A PyTorch loss function to calculate loss on the test data.
    device: A target device to compute on (e.g. "cuda" or "cpu").

    Returns:
    A tuple of testing loss and testing accuracy metrics.
    In the form (test_loss, test_accuracy). For example:

    (0.0223, 0.8985)
    """
    # Put model in eval mode
    model.eval()

    # Setup test loss and test accuracy values
    test_loss, test_acc = 0, 0
    y_pred = []
    y_true = []
    y_proba = []
    softmax = nn.Softmax(dim=1)

    # Turn on inference context manager
    with torch.inference_mode():
        """
        torch.inference_mode is analogous to torch.no_grad : 
        gets better performance by disabling view tracking and version counter bumps
        """
        # Loop through DataLoader batches
        for images, labels in dataloader:
            # Send data to target device
            images, labels = images.to(device), labels.to(device)

            # 1. Forward pass
            output = model(images)

            # 2. Calculate and accumulate probas
            probas_output = softmax(output)
            y_proba.extend(probas_output.detach().cpu().numpy())

            # 3. Calculate and accumulate loss
            loss = loss_fn(output, labels)
            test_loss += loss.item()

            # 4. Calculate and accumulate accuracy
            labels = labels.data.cpu().numpy()
            y_true.extend(labels)  # Save Truth
            preds = np.argmax(output.detach().cpu().numpy(), axis=1)
            y_pred.extend(preds)  # Save Prediction
            acc = (preds == labels).mean()
            test_acc += acc

    y_proba = np.array(y_proba)
    # Adjust metrics to get average loss and accuracy per batch
    test_loss = test_loss / len(dataloader)
    test_acc = test_acc / len(dataloader)
    return test_loss, test_acc * 100, y_pred, y_true, y_proba


def train_step(model: torch.nn.Module, dataloader: torch.utils.data.DataLoader, loss_fn: torch.nn.Module,
               optimizer: torch.optim.Optimizer, device: torch.device) -> Tuple[float, float]:
    """Trains a PyTorch model for a single epoch.

    Turns a target PyTorch model to training mode and then
    runs through all the required training steps (forward
    pass, loss calculation, optimizer step).

    Args:
    model: A PyTorch model to be trained.
    dataloader: A DataLoader instance for the model to be trained on.
    loss_fn: A PyTorch loss function to minimize.
    optimizer: A PyTorch optimizer to help minimize the loss function.
    device: A target device to compute on (e.g. "cuda" or "cpu").

    Returns:
    A tuple of training loss and training accuracy metrics.
    In the form (train_loss, train_accuracy). For example:

    (0.1112, 0.8743)
    """
    # Put model in training mode
    model.train()

    # Setup train loss and train accuracy values
    train_loss, train_acc = 0, 0

    # Loop through data loader data batches
    for batch, (images, labels) in enumerate(dataloader):

        # Send data to target device
        images, labels = images.to(device), labels.to(device)

        # 1. Forward pass
        output = model(images)

        # 2. Calculate  and accumulate loss
        loss = loss_fn(output, labels)
        train_loss += loss.item()

        # 3. Optimizer zero grad
        optimizer.zero_grad()  # Sets the gradients of all optimized torch.Tensor to zero.

        # 4. Loss backward
        loss.backward()

        # 5. Optimizer step
        optimizer.step()

        # Calculate and accumulate accuracy metric across all batches
        y_pred_class = torch.argmax(torch.softmax(output, dim=1), dim=1)
        train_acc += (y_pred_class == labels).sum().item()/len(output)

    # Adjust metrics to get average loss and accuracy per batch
    train_loss = train_loss / len(dataloader)
    train_acc = train_acc / len(dataloader)
    return train_loss, train_acc * 100


def train(model: torch.nn.Module, train_dataloader: torch.utils.data.DataLoader,
          test_dataloader: torch.utils.data.DataLoader, optimizer: torch.optim.Optimizer, loss_fn: torch.nn.Module,
          epochs: int, device: torch.device) -> Dict[str, List]:
    """Trains and tests a PyTorch model.

    Passes a target PyTorch models through train_step() and test()
    functions for a number of epochs, training and testing the model
    in the same epoch loop.

    Calculates, prints and stores evaluation metrics throughout.

    Args:
    model: A PyTorch model to be trained and tested.
    train_dataloader: A DataLoader instance for the model to be trained on.
    test_dataloader: A DataLoader instance for the model to be tested on.
    optimizer: A PyTorch optimizer to help minimize the loss function.
    loss_fn: A PyTorch loss function to calculate loss on both datasets.
    epochs: An integer indicating how many epochs to train for.
    device: A target device to compute on (e.g. "cuda" or "cpu").

    Returns:
    A dictionary of training and testing loss as well as training and
    testing accuracy metrics. Each metric has a value in a list for
    each epoch.
    In the form: {train_loss: [...],
                  train_acc: [...],
                  test_loss: [...],
                  test_acc: [...]}
    """
    # Create empty results dictionary
    results = {"train_loss": [], "train_acc": [], "val_loss": [], "val_acc": []}

    # Loop through training and testing steps for a number of epochs
    for epoch in tqdm(range(epochs), colour="BLUE"):

        train_loss, train_acc = train_step(model=model, dataloader=train_dataloader, loss_fn=loss_fn,
                                           optimizer=optimizer, device=device)

        val_loss, val_acc, _, _, _ = test(model=model, dataloader=test_dataloader, loss_fn=loss_fn, device=device)

        # Print out what's happening
        print(
          f"\tTrain Epoch: {epoch + 1} \t"
          f"Train_loss: {train_loss:.4f} | "
          f"Train_acc: {train_acc:.4f} % | "
          f"Validation_loss: {val_loss:.4f} | "
          f"Validation_acc: {val_acc:.4f} %"
        )

        # Update results dictionary
        results["train_loss"].append(train_loss)
        results["train_acc"].append(train_acc)
        results["val_loss"].append(val_loss)
        results["val_acc"].append(val_acc)

    # Return the filled results at the end of the epochs
    return results

going_modular/engine.py
lcoEntreprise/flower-homomorphic_encryption
import torch
from torch import nn
import torch.nn.functional as F


# Define the architecture
class Net(nn.Module):
    """
    A simple CNN model

    Args:
        num_classes: An integer indicating the number of classes in the dataset.
    """
    def __init__(self, num_classes=10) -> None:
        super(Net, self).__init__()
        # 3 input image channel, 6 output channels, 5x5 square convolution
        self.conv1 = nn.Conv2d(3, 6, 5)
        # Max pooling over a (2, 2) window
        self.pool = nn.MaxPool2d(2, 2)
        # 6 input image channel, 16 output channels, 5x5 square convolution
        self.conv2 = nn.Conv2d(6, 16, 5)
        # an affine operation: y = Wx + b
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, num_classes)  # 1 if num_classes <= 2 else num_classes)

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        """
        Forward pass of the neural network
        """
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        # Flatten the tensor into a vector
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        # output layer
        x = self.fc3(x)
        return x

going_modular/model_builder.py
lcoEntreprise/flower-homomorphic_encryption
from typing import List, Tuple, cast
import numpy as np
import time
import os

# For the homomorphic encryption
import pickle
import tenseal as ts
from flwr.common import NDArray, NDArrays, Parameters
from functools import reduce


# /////////////////////// Homomorphic Encryption \\\\\\\\\\\\\\\\\\\\\\\\\\

def context():
    """
    This function is used to create the context of the homomorphic encryption:
    it is used to create the keys and the parameters of the encryption scheme (CKKS).

    :return: the context of the homomorphic encryption
    """
    cont = ts.context(
        ts.SCHEME_TYPE.CKKS,
        poly_modulus_degree=8192,
        # This means that the coefficient modulus will contain 4 primes of 60 bits, 40 bits, 40 bits, and 60 bits.
        coeff_mod_bit_sizes=[60, 40, 40, 60]
    )

    cont.generate_galois_keys()  # You can create the Galois keys by calling generate_galois_keys
    cont.global_scale = 2 ** 40  # global_scale: the scaling factor, here set to 2**40 (same that pow(2, 40))
    return cont


class Layer:
    """
    This class is used to represent the weights of a layer of a neural network.

    :param name_layer: the name of the layer
    :param weight: the weights of the layer
    """
    def __init__(self, name_layer, weight):
        self.name = name_layer
        self.weight_array = weight

    def get_name(self):
        """
        This function is used to get the name of the layer.
        """
        return self.name

    def get_weight(self):
        """
        This function is used to get the weights of the layer.
        """
        return self.weight_array

    def __add__(self, other):
        """
        This function is used to add the weights of two layers.

        :param other: the other layer object
        """
        weights = other.get_weight() if type(other) == Layer else other
        return Layer(self.name, self.weight_array + weights)

    def __sub__(self, other):
        """
        This function is used to substract the weights of two layers.

        :param other: the other layer object

        :return: the substraction result in the form of a layer object
        """
        weights = other.get_weight() if type(other) == Layer else other
        return Layer(self.name, self.weight_array - weights)

    def __mul__(self, other):
        """
        This function is used to multiply the weights of two layers.

        :param other: the other layer object

        :return: the multiplication result in the form of a layer object
        """
        weights = other.get_weight() if type(other) == Layer else other
        return Layer(self.name, self.weight_array * weights)

    def __truediv__(self, other):
        """
        This function is used to divide the weights of two layers.
        It works only if the weights of the other layer are a number or a layer object (weights with the same shape).

        :param other: the other layer object

        :return: the division result in the form of a layer object
        """
        weights = other.get_weight() if type(other) == Layer else other
        weights = self.weight_array * (1 / weights)
        return Layer(self.name, weights)

    def __len__(self):
        """
        This function is used to get the number of weights of the layer (the size of the layer).
        """
        somme = 1
        for elem in self.weight_array.shape():
            somme *= elem
        return somme

    def shape(self):
        """
        This function is used to get the shape of the layer.
        """
        return self.weight_array.shape()

    def sum(self, axis=0):
        """
        This function is used to get the sum of the weights of the layer.

        :param axis: the axis of the sum (default: 0 for the sum of the columns)

        :return: the sum of the weights of the layer in the form of a layer object
        """
        return Layer(f"sum_{self.name}", self.weight_array.sum(axis=axis))

    def mean(self, axis=0):
        """
        This function is used to get the mean of the weights of the layer.

        :param axis: the axis of the mean (default: 0 for the mean of the columns)

        :return: the mean of the weights of the layer in the form of a layer object
        """
        weights = self.weight_array.sum(axis=axis) * (1 / self.weight_array.shape[axis])
        return Layer(f"sum_{self.name}", weights)

    def decrypt(self, sk=None):
        """
        This function is used to decrypt the weights of the layer.

        :param sk: the secret key used to decrypt the weights (default: None).
        Not used in this class but this is to respect the same structure as the CryptedLayer class.

        :return: the decrypted weights of the layer in the form of a list of weights
        """
        return self.weight_array.tolist()

    def serialize(self):
        """
        This function is used to serialize the weights of the layer.

        :return: the serialized weights of the layer in the form of a dictionary
        with the name of the layer as key and the weights as value
        """

        return {self.name: self.weight_array}


class CryptedLayer(Layer):
    """
    This class is used to represent the crypted weights of a layer of a neural network.

    :param name_layer: the name of the layer
    :param weight: the crypted weights of the layer
    :param contexte: the context of the encryption
    """
    def __init__(self, name_layer, weight, contexte=None):
        super(CryptedLayer, self).__init__(name_layer, weight)
        if type(weight) == ts.tensors.CKKSTensor or type(weight) == bytes:
            # If the weights are already encrypted or if they are bytes (serialized weights)
            self.weight_array = weight

        else:
            # If the weights are not encrypted, we encrypt them with the context
            self.weight_array = ts.ckks_tensor(contexte, weight.cpu().detach().numpy())

    def __add__(self, other):
        """
        This function is used to add the crypted weights of two layers.

        :param other: the other layer object

        :return: the addition result in the form of a crypted layer object
        """
        weights = other.get_weight() if type(other) == CryptedLayer else other
        return CryptedLayer(self.name, self.weight_array + weights)

    def __sub__(self, other):
        """
        This function is used to substract the crypted weights of two layers.

        :param other: the other layer object

        :return: the substraction result in the form of a crypted layer object
        """
        weights = other.get_weight() if type(other) == CryptedLayer else other
        return CryptedLayer(self.name, self.weight_array - weights)

    def __mul__(self, other):
        """
        This function is used to multiply the crypted weights of two layers.

        :param other: the other layer object

        :return: the multiplication result in the form of a crypted layer object
        """
        weights = other.get_weight() if type(other) == CryptedLayer else other
        return CryptedLayer(self.name, self.weight_array * weights)

    def __truediv__(self, other):
        """
        This function is used to divide the crypted weights of two layers.

        :param other: the other layer object

        :return: the division result in the form of a crypted layer object
        """
        try:
            # We try to divide the weights of the layer by the weights of the other layer or by a number
            # It is possible only if the denominator is a number or a tensor of non-crypted weights
            weights = other.get_weight() if type(other) == CryptedLayer else other
            weights = self.weight_array * (1 / weights)

        except:
            print("Error: the division operator isn't supported by SEAL")
            weights = []

        return CryptedLayer(self.name, weights)

    def shape(self):
        """
        This function is used to get the shape of the layer.

        :return: the shape of the layer in the form of a tuple
        """
        return self.weight_array.shape

    def sum(self, axis=0):
        """
        This function is used to get the sum of the weights of the layer.

        :param axis: the axis of the sum (default: 0 for the sum of the columns)

        :return: the sum of the weights of the layer in the form of a crypted layer object
        """
        return CryptedLayer(f"sum_{self.name}", self.weight_array.sum(axis=axis))

    def mean(self, axis=0):
        """
        This function is used to get the mean of the weights of the layer.

        :param axis: the axis of the mean (default: 0 for the mean of the columns)

        :return: the average of the weights of the layer in the form of a crypted layer object
        """
        weights = self.weight_array.sum(axis=axis) * (1 / self.weight_array.shape[axis])
        return CryptedLayer(f"sum_{self.name}", weights)

    def decrypt(self, sk=None):
        """
        This function is used to decrypt the weights of the layer.

        :param sk: the secret key used to decrypt the weights (default: None)

        :return: the decrypted weights of the layer in the form of a list of weights
        """
        return self.weight_array.decrypt(sk).tolist() if sk else self.weight_array.decrypt().tolist()

    def serialize(self):
        """
        This function is used to serialize the weights of the layer.

        :return: the serialized weights of the layer in the form of a dictionary
        """
        return {self.name: self.weight_array.serialize()}


def crypte(client_w, context_c):
    """
    This function is used to crypte the weights of a neural network.

    :param client_w: dictionary of the weights of a neural network
    :param context_c: the context of the encryption
    :return: a list of Layer objects (crypted weights or not)
    """
    encrypted = []
    for name_layer, weight_array in client_w.items():
        start_time = time.time()
        if name_layer == 'fc3.weight':
            encrypted.append(CryptedLayer(name_layer, weight_array, context_c))

        else:
            encrypted.append(Layer(name_layer, weight_array))

        print(name_layer, (time.time() - start_time))

    # return [CryptedLayer(name_layer, weight_array, context_c) for name_layer, weight_array in client_w.items()]
    return encrypted


def read_query(file_path):
    """
    This function is used to read a pickle file.

    :param file_path: the path of the file to read
    :return: the query and the context
    """
    if os.path.exists(file_path):
        with open(file_path, 'rb') as file:
            """
            # pickle.load(f)  # load to read file object

            file_str = f.read()
            client_query1 = pickle.loads(file_str)  # loads to read str class
            """
            query_str = pickle.load(file)

        contexte = query_str["contexte"]  # ts.context_from(query["contexte"])
        del query_str["contexte"]
        return query_str, contexte

    else:
        print("The file doesn't exist")


def write_query(file_path, client_query):
    """
    This function is used to write a pickle file.

    :param file_path: the path of the file to write

    :param client_query: the query to write
    """
    with open(file_path, 'wb') as file:  # 'ab' to add existing file
        encode_str = pickle.dumps(client_query)
        file.write(encode_str)


def deserialized_layer(name_layer, weight_array, ctx):
    """
    This function is used to deserialized a layer (crypted or not).
    i.e. to transform the weights of a layer in the correct format (not bytes).

    :param name_layer: the name of the layer
    :param weight_array: the weights of the layer
    :param ctx: the context (if the layer is crypted)
    :return: the object Layer or CryptedLayer with the weights of the layer in the correct format
    """
    if type(weight_array) == bytes:
        return CryptedLayer(name_layer, ts.ckks_tensor_from(ctx, weight_array), ctx)

    elif type(weight_array) == ts.tensors.CKKSTensor:
        return CryptedLayer(name_layer, weight_array, ctx)

    else:
        return Layer(name_layer, weight_array)


def deserialized_model(client_query, ctx):
    """
    This function is used to deserialized a model (crypted or not).
    i.e. to transform the weights of a model in the correct format (not bytes).

    :param client_query: the model
    :param ctx: the context (if the model is crypted)
    :return: the model with the weights in the correct format
    """
    return [deserialized_layer(name_layer, weight_array, ctx) for name_layer, weight_array in client_query.items()]


# Redefine the aggregate function (defined in Flower)
def aggregate_custom(results: List[Tuple[NDArrays, int]]) -> NDArrays:
    """Compute weighted average.
    Args:
        results: List of tuples containing the model weights and the number of samples used to compute the weights.
    Returns: A list of model weights averaged according to the number of samples used to compute the weights.
    """
    # Calculate the total number of examples used during training
    num_examples_total = sum([num_examples for _, num_examples in results])
    # Create a list of weights, each multiplied by the related number of examples
    weighted_weights = [
        [layer * num_examples for layer in weights] for weights, num_examples in results
    ]
    # Compute average weights of each layer
    weights_prime: NDArrays = [
        reduce(np.add, layer_updates) * (1/num_examples_total)
        for layer_updates in zip(*weighted_weights)
    ]
    return weights_prime

going_modular/security.py
lcoEntreprise/flower-homomorphic_encryption
from client import *
from federated import *
"""
script to start the client with the federated learning pipeline with Flower.
"""

if __name__ == '__main__':
    print("start client")
    start_time = time.time()
    fl.client.start_numpy_client(server_address="[::]:8080",
                                 client=client_common(args.id_client, args.model_save, args.yaml_path, args.roc_path,
                                                      args.save_results, args.matrix_path,
                                                      args.batch_size, trainloaders, valloaders,
                                                      DEVICE, CLASSES, args.he, args.path_keys, args.path_crypted)
                                 )
    print(f"client Time = {time.time() - start_time} seconds")

main_client.py
lcoEntreprise/flower-homomorphic_encryption
from server import *
"""
Script to start the server side of the federated learning pipeline with Flower.
"""

if __name__ == '__main__':
    print("start server")
    start_time = time.time()

    fl.server.start_server(server_address="0.0.0.0:8080",
                           config=fl.server.ServerConfig(num_rounds=args.rounds),
                           strategy=strategy)
    print(f"Server Time = {time.time() - start_time} seconds")

main_server.py
lcoEntreprise/flower-homomorphic_encryption
import torch.nn.functional
from typing import Optional, Callable, Union
from logging import WARNING
import flwr.server.strategy
from flwr.server.client_proxy import ClientProxy
from flwr.server.client_manager import ClientManager
from flwr.server.strategy.aggregate import weighted_loss_avg
from flwr.common import (
    Metrics, EvaluateIns, EvaluateRes, FitIns, FitRes, MetricsAggregationFn, Scalar, logger,
    ndarrays_to_parameters_custom, parameters_to_ndarrays_custom
)

from federated import *
from going_modular.security import aggregate_custom

"""
Script with the server-side logic for federated learning with Flower.
"""
# #############################################################################
# 3. Federated learning strategy
# #############################################################################
central = Net(num_classes=len(CLASSES)).to(DEVICE)


# Weighted averaging function to aggregate the accuracy metric we return from evaluate
def weighted_average(metrics: List[Tuple[int, Metrics]]) -> Metrics:
    # Multiply accuracy of each client by number of examples used
    accuracies = [num_examples * m["accuracy"] for num_examples, m in metrics]
    examples = [num_examples for num_examples, _ in metrics]

    # Aggregate and return custom metric (weighted average)
    return {"accuracy": sum(accuracies) / sum(examples)}


# Function will be by Flower called after every round : Centralized Evaluation (or server-side evaluation)
def evaluate2(server_round: int, parameters: NDArrays,
              config: Dict[str, Scalar]) -> Optional[Tuple[float, Dict[str, Scalar]]]:
    set_parameters(central, parameters)  # Update model with the latest parameters
    loss, accuracy, y_pred, y_true, y_proba = engine.test(central, testloader, loss_fn=torch.nn.CrossEntropyLoss(),
                                                          device=DEVICE)
    print(f"Server-side evaluation loss {loss} / accuracy {accuracy}")
    return loss, {"accuracy": accuracy}


def get_on_fit_config_fn(epoch=2, lr=0.001, batch_size=32) -> Callable[[int], Dict[str, str]]:
    """Return a function which returns training configurations."""

    def fit_config(server_round: int) -> Dict[str, str]:
        """
        Return training configuration dict for each round with static batch size and (local) epochs.

        Perform two rounds of training with one local epoch, increase to two local epochs afterwards.
        """
        config = {
            "learning_rate": str(lr),
            "batch_size": str(batch_size),
            "server_round": server_round,  # The current round of federated learning
            "local_epochs": epoch  # 1 if server_round < 2 else epoch,
        }
        return config

    return fit_config


def aggreg_fit_checkpoint(server_round, aggregated_parameters, central_model, path_checkpoint,
                          context_client=None, server_path=""):
    if aggregated_parameters is not None:
        print(f"Saving round {server_round} aggregated_parameters...")

        # Convert `Parameters` to `List[np.ndarray]`
        aggregated_ndarrays: List[np.ndarray] = parameters_to_ndarrays_custom(aggregated_parameters, context_client)
        """
        [value of 'conv1.weight', 
        value of 'conv1.bias', 
        value of 'conv2.weight', 
        value of 'conv2.bias', 
        value of 'fc1.weight', 
        value of 'fc1.bias', 
        value of 'fc2.weight', 
        value of 'fc2.bias', 
        value of 'fc3.weight', 
        value of 'fc3.bias']
        """
        if context_client:
            def serialized(key, matrix):
                if key == 'fc3.weight':
                    return matrix.serialize()
                else:
                    return matrix

            server_response = {
                **{
                    key: serialized(key, aggregated_ndarrays[i]) for i, key in enumerate(central_model.state_dict().keys())},  # bytes
                "contexte": server_context.serialize()
            }

            security.write_query(server_path, server_response)

        else:
            # Convert `List[np.ndarray]` to PyTorch`state_dict`
            params_dict = zip(central_model.state_dict().keys(), aggregated_ndarrays)
            state_dict = OrderedDict({k: torch.tensor(v) for k, v in params_dict})
            central_model.load_state_dict(state_dict, strict=True)

            # Save the model
            if path_checkpoint:
                torch.save({
                    'model_state_dict': central_model.state_dict(),
                }, path_checkpoint)


# A Strategy from scratch with the same sampling of the clients as it is in FedAvg
# and then change the configuration dictionary
class FedCustom(fl.server.strategy.Strategy):
    def __init__(
            self,
            fraction_fit: float = 1.0,
            fraction_evaluate: float = 1.0,
            min_fit_clients: int = 2,
            min_evaluate_clients: int = 2,
            min_available_clients: int = 2,
            evaluate_fn: Optional[
                    Callable[[int, NDArrays, Dict[str, Scalar]], Optional[Tuple[float, Dict[str, Scalar]]]]
                ] = None,
            on_fit_config_fn: Optional[Callable[[int], Dict[str, Scalar]]] = None,
            on_evaluate_config_fn: Optional[Callable[[int], Dict[str, Scalar]]] = None,
            accept_failures: bool = True,
            initial_parameters: Optional[Parameters] = None,
            fit_metrics_aggregation_fn: Optional[MetricsAggregationFn] = None,
            evaluate_metrics_aggregation_fn: Optional[MetricsAggregationFn] = None,
            context_client=None
    ) -> None:
        super().__init__()
        self.fraction_fit = fraction_fit
        self.fraction_evaluate = fraction_evaluate
        self.min_fit_clients = min_fit_clients
        self.min_evaluate_clients = min_evaluate_clients
        self.min_available_clients = min_available_clients
        self.evaluate_fn = evaluate_fn
        self.on_fit_config_fn = on_fit_config_fn
        self.on_evaluate_config_fn = on_evaluate_config_fn,
        self.accept_failures = accept_failures
        self.initial_parameters = initial_parameters
        self.fit_metrics_aggregation_fn = fit_metrics_aggregation_fn
        self.evaluate_metrics_aggregation_fn = evaluate_metrics_aggregation_fn
        self.context_client = context_client

    def __repr__(self) -> str:
        # Same function as FedAvg(Strategy)
        return f"FedCustom (accept_failures={self.accept_failures})"

    def initialize_parameters(
        self, client_manager: ClientManager
    ) -> Optional[Parameters]:
        """Initialize global model parameters."""
        # Same function as FedAvg(Strategy)
        initial_parameters = self.initial_parameters
        self.initial_parameters = None  # Don't keep initial parameters in memory
        return initial_parameters

    def num_fit_clients(self, num_available_clients: int) -> Tuple[int, int]:
        """Return sample size and required number of clients."""
        # Same function as FedAvg(Strategy)
        num_clients = int(num_available_clients * self.fraction_fit)
        return max(num_clients, self.min_fit_clients), self.min_available_clients

    def configure_fit(
        self, server_round: int, parameters: Parameters, client_manager: ClientManager
    ) -> List[Tuple[ClientProxy, FitIns]]:
        """Configure the next round of training."""
        # Sample clients
        sample_size, min_num_clients = self.num_fit_clients(
            client_manager.num_available()
        )

        clients = client_manager.sample(
            num_clients=sample_size, min_num_clients=min_num_clients
        )
        # Create custom configs
        n_clients = len(clients)
        half_clients = n_clients // 2
        # Custom fit config function provided
        standard_lr = args.lr
        higher_lr = 0.003
        config = {"server_round": server_round, "local_epochs": 1}
        if self.on_fit_config_fn is not None:
            # Custom fit config function provided
            config = self.on_fit_config_fn(server_round)

        # fit_ins = FitIns(parameters, config)
        # Return client/config pairs
        fit_configurations = []
        for idx, client in enumerate(clients):
            config["learning_rate"] = standard_lr if idx < half_clients else higher_lr
            """
            Each pair of (ClientProxy, FitRes) constitutes 
            a successful update from one of the previously selected clients.
            """
            fit_configurations.append(
                (
                    client,
                    FitIns(
                        parameters,
                        config
                    )
                )
            )
        # Successful updates from the previously selected and configured clients
        return fit_configurations

    def aggregate_fit(
        self,
        server_round: int,
        results: List[Tuple[ClientProxy, FitRes]],
        failures: List[Union[Tuple[ClientProxy, FitRes], BaseException]],
    ) -> Tuple[Optional[Parameters], Dict[str, Scalar]]:
        """Aggregate fit results using weighted average. (each round)"""
        # Same function as FedAvg(Strategy)
        if not results:
            return None, {}

        # Do not aggregate if there are failures and failures are not accepted
        if not self.accept_failures and failures:
            return None, {}

        # Convert results parameters --> array matrix
        weights_results = [
            (parameters_to_ndarrays_custom(fit_res.parameters, self.context_client), fit_res.num_examples)
            for _, fit_res in results
        ]

        # Aggregate parameters using weighted average between the clients and convert back to parameters object (bytes)
        parameters_aggregated = ndarrays_to_parameters_custom(aggregate_custom(weights_results))

        metrics_aggregated = {}
        # Aggregate custom metrics if aggregation fn was provided
        if self.fit_metrics_aggregation_fn:
            fit_metrics = [(res.num_examples, res.metrics) for _, res in results]
            metrics_aggregated = self.fit_metrics_aggregation_fn(fit_metrics)

        elif server_round == 1:  # Only log this warning once
            logger.log(WARNING, "No fit_metrics_aggregation_fn provided")

        # Same function as SaveModelStrategy(fl.server.strategy.FedAvg)
        """Aggregate model weights using weighted average and store checkpoint"""
        aggreg_fit_checkpoint(server_round, parameters_aggregated, central, args.model_save,
                              self.context_client, args.path_crypted)
        return parameters_aggregated, metrics_aggregated

    def num_evaluation_clients(self, num_available_clients: int) -> Tuple[int, int]:
        """Use a fraction of available clients for evaluation."""
        # Same function as FedAvg(Strategy)
        num_clients = int(num_available_clients * self.fraction_evaluate)
        return max(num_clients, self.min_evaluate_clients), self.min_available_clients

    def configure_evaluate(
        self, server_round: int, parameters: Parameters, client_manager: ClientManager
    ) -> List[Tuple[ClientProxy, EvaluateIns]]:
        """Configure the next round of evaluation."""
        # Same function as FedAvg(Strategy)
        # Do not configure federated evaluation if fraction eval is 0.
        if self.fraction_evaluate == 0.0:
            return []

        # Parameters and config
        config = {}  # {"server_round": server_round, "local_epochs": 1}

        evaluate_ins = EvaluateIns(parameters, config)

        # Sample clients
        sample_size, min_num_clients = self.num_evaluation_clients(
            client_manager.num_available()
        )

        clients = client_manager.sample(
            num_clients=sample_size, min_num_clients=min_num_clients
        )

        # Return client/config pairs
        # Each pair of (ClientProxy, FitRes) constitutes a successful update from one of the previously selected clients
        return [(client, evaluate_ins) for client in clients]

    def aggregate_evaluate(
        self,
        server_round: int,
        results: List[Tuple[ClientProxy, EvaluateRes]],
        failures: List[Union[Tuple[ClientProxy, EvaluateRes], BaseException]],
    ) -> Tuple[Optional[float], Dict[str, Scalar]]:
        """Aggregate evaluation losses using weighted average."""
        # Same function as FedAvg(Strategy)
        if not results:
            return None, {}

        # Do not aggregate if there are failures and failures are not accepted
        if not self.accept_failures and failures:
            return None, {}

        # Aggregate loss
        loss_aggregated = weighted_loss_avg(
            [
                (evaluate_res.num_examples, evaluate_res.loss)
                for _, evaluate_res in results
            ]
        )

        metrics_aggregated = {}
        # Aggregate custom metrics if aggregation fn was provided
        if self.evaluate_metrics_aggregation_fn:
            eval_metrics = [(res.num_examples, res.metrics) for _, res in results]
            metrics_aggregated = self.evaluate_metrics_aggregation_fn(eval_metrics)

        # Only log this warning once
        elif server_round == 1:
            logger.log(WARNING, "No evaluate_metrics_aggregation_fn provided")

        return loss_aggregated, metrics_aggregated

    def evaluate(
        self, server_round: int, parameters: Parameters
    ) -> Optional[Tuple[float, Dict[str, Scalar]]]:
        """Evaluate global model parameters using an evaluation function."""
        # Same function as FedAvg(Strategy)
        if self.evaluate_fn is None:
            # Let's assume we won't perform the global model evaluation on the server side.
            return None

        # if we have a global model evaluation on the server side :
        parameters_ndarrays = parameters_to_ndarrays_custom(parameters, self.context_client)
        eval_res = self.evaluate_fn(server_round, parameters_ndarrays, {})

        # if you haven't results
        if eval_res is None:
            return None

        loss, metrics = eval_res
        return loss, metrics


"""
- Setting `min_available_clients` lower than `min_fit_clients` or
`min_evaluate_clients` can cause the server to fail when there are too few clients
connected to the server. 

- `min_available_clients` must be set to a value larger
than or equal to the values of `min_fit_clients` and `min_evaluate_clients`.
"""
if args.he:
    print("get public key : ", args.path_public_key)
    _, server_context = security.read_query(args.path_public_key)
    server_context = ts.context_from(server_context)

else:
    print("not HE so not key")
    server_context = None

strategy = FedCustom(
    fraction_fit=args.frac_fit,  # Train on frac_fit % clients (each round)
    fraction_evaluate=args.frac_eval,  # Sample frac_eval % of available clients for evaluation
    min_fit_clients=args.min_fit_clients,  # Never sample less than 10 clients for training
    min_evaluate_clients=args.min_eval_clients if args.min_eval_clients else args.number_clients // 2,
    min_available_clients=args.min_avail_clients,  # Wait until all 10 clients are available
    evaluate_metrics_aggregation_fn=weighted_average,  # <-- pass the metric aggregation function
    initial_parameters=ndarrays_to_parameters_custom(get_parameters2(central)),  # prevents Flower from asking one of the clients for the initial parameters
    evaluate_fn=None if args.he else evaluate2,  # Pass the evaluation function for the server side
    on_fit_config_fn=get_on_fit_config_fn(epoch=args.max_epochs, lr=args.lr, batch_size=args.batch_size),  # Pass the fit_config function
    context_client=server_context,
)


server.py
lcoEntreprise/flower-homomorphic_encryption
from client import *
from server import *

"""
Script to start the simulation of the federated learning pipeline with Flower (client and server).
"""


# Implementing a Flower client
def client_fn(cid: str) -> FlowerClient:
    """Create a Flower client representing a single organization."""
    return client_common(cid,
                         model_save, path_yaml, path_roc, results_save, path_matrix,
                         batch_size, trainloaders, valloaders, DEVICE, CLASSES, he, secret_path, server_path)


# ////////////////////////////// Simulation of the federated learning pipeline with Flower ////////////////////////////
# Pass parameters to the Strategy for server-side parameter initialization
if __name__ == '__main__':
    # Specify client resources if you need GPU (defaults to 1 CPU and 0 GPU)
    """
    client_resources = None
    if DEVICE.type == "cuda":
        client_resources = {"num_gpus": 1}
    """
    model_save = args.model_save
    path_yaml = args.yaml_path
    path_roc = args.roc_path
    results_save = args.save_results
    path_matrix = args.matrix_path
    batch_size = args.batch_size
    he = args.he
    secret_path = args.path_keys
    server_path = args.path_crypted

    print("Start simulation")
    start_simulation = time.time()
    fl.simulation.start_simulation(
        client_fn=client_fn,
        num_clients=args.number_clients,
        config=fl.server.ServerConfig(num_rounds=args.rounds),
        strategy=strategy,
        # client_resources=client_resources
    )
    print(f"Simulation Time = {time.time() - start_simulation} seconds")

simulation.py
lcoEntreprise/flower-homomorphic_encryption