code: # -*- coding: utf-8 -*-
  """
  Functions for model training and evaluation (single-partner and multi-partner cases)
  """
  
  import operator
  import os
  from abc import ABC, abstractmethod
  from copy import deepcopy
  from timeit import default_timer as timer
  
  import numpy as np
  import random
  import tensorflow as tf
  from loguru import logger
  from sklearn.metrics import confusion_matrix
  from tensorflow.keras import Input, Model
  from tensorflow.keras.backend import clear_session
  from tensorflow.keras.callbacks import EarlyStopping
  
  from .utils import History
  from ..utils import project_onto_the_simplex
  from .. import constants
  from ..models import NoiseAdaptationChannel, EnsemblePredictionsModel
  from ..partner import Partner, PartnerMpl
  
  ALLOWED_PARAMETERS = ('partners_list',
                        'epoch_count',
                        'minibatch_count',
                        'dataset',
                        'aggregation',
                        'is_early_stopping',
                        'is_save_data',
                        'save_folder',
                        'init_model_from',
                        'use_saved_weights')
  
  
  class MultiPartnerLearning(ABC):
      name = 'abstract'
  
      def __init__(self, scenario, **kwargs):
          """
  
          :type scenario: Scenario
          """
          # Attributes related to the data and the model
          self.dataset = scenario.dataset
          self.partners_list = scenario.partners_list
          self.init_model_from = scenario.init_model_from
          self.use_saved_weights = scenario.use_saved_weights
          self.amounts_per_partner = scenario.amounts_per_partner
          self.val_set = scenario.val_set
          self.test_set = scenario.test_set
  
          # Attributes related to iterating at different levels
          self.epoch_count = scenario.epoch_count
          self.minibatch_count = scenario.minibatch_count
          self.is_early_stopping = scenario.is_early_stopping
  
          # Attributes to store results
          self.save_folder = scenario.save_folder
  
          # Erase the default parameters (which mostly come from the scenario) if some parameters have been specified
          self.__dict__.update((k, v) for k, v in kwargs.items() if k in ALLOWED_PARAMETERS)
  
          # Unpack dataset-related parameters
          self.val_data = (self.dataset.x_val, self.dataset.y_val)
          self.test_data = (self.dataset.x_test, self.dataset.y_test)
          self.dataset_name = self.dataset.name
          self.generate_new_model = self.dataset.generate_new_model
  
          # Initialize the model
          model = self.init_model()
          self.model_weights = model.get_weights()
          self.metrics_names = self.dataset.model_metrics_names
  
          # Initialize iterators
          self.epoch_index = 0
          self.minibatch_index = 0
          self.learning_computation_time = 0
  
          # Convert partners to Mpl partners
          for partner in self.partners_list:
              assert isinstance(partner, Partner)
          partners_list = sorted(self.partners_list, key=operator.attrgetter("id"))
          logger.info(
              f"## Preparation of model's training on partners with ids: {['#' + str(p.id) for p in partners_list]}")
          self.partners_list = [PartnerMpl(partner, self) for partner in self.partners_list]
  
          # Attributes related to the aggregation approach
          self.aggregator = self.init_aggregation_function(scenario.aggregation)
  
          # Initialize History
          self.history = History(self)
  
          # Initialize result folder
          if self.save_folder is not None:
              if 'custom_name' in kwargs:
                  self.save_folder = self.save_folder / kwargs["custom_name"]
              else:
                  self.save_folder = self.save_folder / 'multi_partner_learning'
                  self.save_folder.mkdir(parents=True, exist_ok=False)
  
          logger.debug("MultiPartnerLearning object instantiated.")
  
      def __str__(self):
          return f'{self.name}'
  
      @property
      def partners_count(self):
          return len(self.partners_list)
  
      def init_aggregation_function(self, aggregator):
          return aggregator(self)
  
      def build_model(self):
          return self.build_model_from_weights(self.model_weights)
  
      def build_model_from_weights(self, new_weights):
          """Generate a new model initialized with weights passed as arguments"""
          new_model = self.generate_new_model()
          new_model.set_weights(new_weights)
          return new_model
  
      def init_model(self):
          new_model = self.generate_new_model()
  
          if self.use_saved_weights:
              logger.info("Init model with previous coalition model")
              new_model.load_weights(self.init_model_from)
          else:
              logger.info("Init new model")
  
          return new_model
  
      def save_final_model(self):
          """Save final model weights"""
  
          model_folder = os.path.join(self.save_folder, 'model')
  
          if not os.path.isdir(model_folder):
              os.makedirs(model_folder)
  
          np.save(os.path.join(model_folder, self.dataset_name + '_final_weights.npy'), self.model_weights)
  
          model_to_save = self.build_model()
          model_to_save.save_weights(os.path.join(model_folder, self.dataset_name + '_final_weights.h5'))
  
      def save_data(self):
          if self.save_folder is None:
              raise ValueError("The path to the save folder is None, history data cannot be saved, nor model weights")
  
          self.save_final_model()
          self.history.save_data()
  
      def log_partner_perf(self, partner_id, partner_index, history):
          for key_history in self.history.metrics:
              self.history.history[partner_id][key_history][self.epoch_index,
                                                            self.minibatch_index] = history[key_history][-1]
  
          epoch_nb_str = f"Epoch {str(self.epoch_index).zfill(2)}/{str(self.epoch_count - 1).zfill(2)}"
          mb_nb_str = f"Minibatch {str(self.minibatch_index).zfill(2)}/{str(self.minibatch_count - 1).zfill(2)}"
          partner_id_str = f"Partner partner_id #{partner_id} ({partner_index}/{self.partners_count - 1})"
          val_acc_str = f"{round(history['val_accuracy'][-1], 2)}"
  
          logger.debug(f"{epoch_nb_str} > {mb_nb_str} > {partner_id_str} > val_acc: {val_acc_str}")
  
      def eval_and_log_model_val_perf(self):
  
          model = self.build_model()
  
          if self.val_set == 'global':
              hist = model.evaluate(self.val_data[0],
                                    self.val_data[1],
                                    batch_size=constants.DEFAULT_BATCH_SIZE,
                                    verbose=0,
                                    )
          elif self.val_set == 'local':
              hist = [0.0, 0.0]
              for p in self.partners_list:
                  hist_partner = model.evaluate(p.x_val,
                                                p.y_val,
                                                batch_size=constants.DEFAULT_BATCH_SIZE,
                                                verbose=0,
                                                )
                  hist[0] += hist_partner[0] / self.partners_count
                  hist[1] += hist_partner[1] / self.partners_count
          else:
              raise ValueError("validation set should be 'local' or 'global', not {self.val_set}")
  
          self.history.history['mpl_model']['val_loss'][self.epoch_index, self.minibatch_index] = hist[0]
          self.history.history['mpl_model']['val_accuracy'][self.epoch_index, self.minibatch_index] = hist[1]
  
          if self.minibatch_index >= self.minibatch_count - 1:
              epoch_nb_str = f"{str(self.epoch_index).zfill(2)}/{str(self.epoch_count - 1).zfill(2)}"
              logger.info(f"   Model evaluation at the end of the epoch "
                          f"{epoch_nb_str}: "
                          f"{['%.3f' % elem for elem in hist]}")
  
      def eval_and_log_final_model_test_perf(self):
          logger.info("### Evaluating model on test data:")
          model = self.build_model()
          if self.test_set == 'global':
              hist = model.evaluate(self.test_data[0],
                                    self.test_data[1],
                                    batch_size=constants.DEFAULT_BATCH_SIZE,
                                    verbose=0,
                                    )
          elif self.test_set == 'local':
              hist = [0.0, 0.0]
              for p in self.partners_list:
                  hist_partner = model.evaluate(p.x_test,
                                                p.y_test,
                                                batch_size=constants.DEFAULT_BATCH_SIZE,
                                                verbose=0,
                                                )
                  hist[0] += hist_partner[0] / self.partners_count
                  hist[1] += hist_partner[1] / self.partners_count
          else:
              raise ValueError("test set should be 'local' or 'global', not {self.val_set}")
  
          self.history.score = hist[1]
          self.history.nb_epochs_done = self.epoch_index + 1
          logger.info(f"   Model metrics names: {self.metrics_names}")
          logger.info(f"   Model metrics values: {['%.3f' % elem for elem in hist]}")
  
      def split_in_minibatches(self):
          """Split the dataset passed as argument in mini-batches"""
  
          for partner in self.partners_list:
              partner.split_minibatches()
  
      def early_stop(self):
          logger.debug("      Checking if early stopping criteria are met:")
          if self.is_early_stopping:
              # Early stopping parameters
              if (
                      self.epoch_index >= constants.PATIENCE
                      and self.history.history['mpl_model']['val_loss'][self.epoch_index,
                                                                        self.minibatch_index] >
                      self.history.history['mpl_model']['val_loss'][self.epoch_index - constants.PATIENCE,
                                                                    self.minibatch_index]
              ):
                  logger.debug("         -> Early stopping criteria are met, stopping here.")
                  return True
              else:
                  logger.debug("         -> Early stopping criteria are not met, continuing with training.")
          else:
              return False
  
      def fit(self):
          """Return the score on test data of a final aggregated model trained in a federated way on each partner"""
  
          start = timer()
          # Train model (iterate for each epoch and mini-batch)
          while self.epoch_index < self.epoch_count:
  
              self.fit_epoch()  # perform an epoch on the self.model
  
              if self.early_stop():
                  break
              self.epoch_index += 1
  
          # After last epoch or if early stopping was triggered, evaluate model on the global testset
          self.eval_and_log_final_model_test_perf()
  
          end = timer()
          self.learning_computation_time = end - start
          logger.info(f"Training and evaluation on multiple partners: "
                      f"done. ({np.round(self.learning_computation_time, 3)} seconds)")
          if self.save_folder is not None:
              self.save_data()  # Save the model weights and the history data
  
      @abstractmethod
      def fit_epoch(self):
          while self.minibatch_index < self.minibatch_count:
              self.fit_minibatch()
              self.minibatch_index += 1
              self.eval_and_log_model_val_perf()
  
      @abstractmethod
      def fit_minibatch(self):
          pass
  
  
  class SinglePartnerLearning(MultiPartnerLearning):
      name = 'Single Partner learning'
  
      def __init__(self, scenario, **kwargs):
          super(SinglePartnerLearning, self).__init__(scenario, **kwargs)
          if self.partners_count != 1:
              raise ValueError('More than one partner is provided')
          self.partner = self.partners_list[0]
  
      def fit(self):
          """Return the score on test data of a model trained on a single partner"""
  
          start = timer()
          logger.info(f"## Training and evaluating model on partner with partner_id #{self.partner.id}")
  
          # Set if early stopping if needed
          cb = []
          es = None
          if self.is_early_stopping:
              es = EarlyStopping(monitor='val_loss', mode='min', verbose=0, patience=constants.PATIENCE)
              cb.append(es)
  
          # Train model
          logger.info("   Training model...")
          model = self.build_model()
          if self.val_set == 'global':
              history = model.fit(self.partner.x_train,
                                  self.partner.y_train,
                                  batch_size=self.partner.batch_size,
                                  epochs=self.epoch_count,
                                  verbose=0,
                                  validation_data=self.val_data,
                                  callbacks=cb)
          elif self.val_set == 'local':
              history = model.fit(self.partner.x_train,
                                  self.partner.y_train,
                                  batch_size=self.partner.batch_size,
                                  epochs=self.epoch_count,
                                  verbose=0,
                                  validation_data=(self.partner.x_val, self.partner.y_val),
                                  callbacks=cb)
          else:
              raise ValueError("validation set should be 'local' or 'global', not {self.val_set}")
  
          self.model_weights = model.get_weights()
          self.log_partner_perf(self.partner.id, 0, history.history)
          del self.history.history['mpl_model']
          # Evaluate trained model on test data
          self.eval_and_log_final_model_test_perf()
          self.history.nb_epochs_done = (es.stopped_epoch + 1) if es.stopped_epoch != 0 else self.epoch_count
  
          end = timer()
          self.learning_computation_time = end - start
  
      def fit_epoch(self):
          pass
  
      def fit_minibatch(self):
          pass
  
  
  class FederatedAverageLearning(MultiPartnerLearning):
      name = 'Federated averaging'
  
      def __init__(self, scenario, **kwargs):
          # First, if only one partner, fall back to dedicated single partner function
          super(FederatedAverageLearning, self).__init__(scenario, **kwargs)
          if self.partners_count == 1:
              raise ValueError('Only one partner is provided. Please use the dedicated SinglePartnerLearning class')
  
      def fit_epoch(self):
          # Clear Keras' old models
          clear_session()
  
          # Split the train dataset in mini-batches
          self.split_in_minibatches()
  
          # Iterate over mini-batches and train
          for i in range(self.minibatch_count):
              self.minibatch_index = i
              self.fit_minibatch()
  
              # At the end of each minibatch,aggregate the models
              self.model_weights = self.aggregator.aggregate_model_weights()
          self.minibatch_index = 0
  
      def fit_minibatch(self):
          """Proceed to a collaborative round with a federated averaging approach"""
  
          logger.debug("Start new fedavg collaborative round ...")
  
          # Starting model for each partner is the aggregated model from the previous mini-batch iteration
          logger.info(f"(fedavg) Minibatch n°{self.minibatch_index} of epoch n°{self.epoch_index}, "
                      f"init each partner's models with a copy of the global model")
  
          for partner in self.partners_list:
              partner.model_weights = self.model_weights
  
          # Evaluate and store accuracy of mini-batch start model
          self.eval_and_log_model_val_perf()
  
          # Iterate over partners for training each individual model
          for partner_index, partner in enumerate(self.partners_list):
              # Reference the partner's model
              partner_model = partner.build_model()
  
              # Train on partner local data set
              if self.val_set == 'global':
                  history = partner_model.fit(partner.minibatched_x_train[self.minibatch_index],
                                              partner.minibatched_y_train[self.minibatch_index],
                                              batch_size=partner.batch_size,
                                              verbose=0,
                                              validation_data=self.val_data)
              elif self.val_set == 'local':
                  history = partner_model.fit(partner.minibatched_x_train[self.minibatch_index],
                                              partner.minibatched_y_train[self.minibatch_index],
                                              batch_size=partner.batch_size,
                                              verbose=0,
                                              validation_data=(partner.x_val, partner.y_val))
              else:
                  raise ValueError("validation set should be 'local' or 'global', not {self.val_set}")
  
              # Log results of the round
              self.log_partner_perf(partner.id, partner_index, history.history)
  
              # Update the partner's model in the models' list
              partner.model_weights = partner_model.get_weights()
  
          logger.debug("End of fedavg collaborative round.")
  
  
  class DistributionallyRobustFederatedAveragingLearning(MultiPartnerLearning):
      """
       - This class implements the Distributionally Robust Federated Averaging (DRFA) Algorithm,
        only a subset of partners are chosen to participate in a given collaborative
       learning round. based on a global mixing parameter called lambda
       - Lambda is updated at the end of each collaborative learning round using its own update rule
       - DRFA is considered a framework under which we can implement other FL algorithms such as FedAvg
       - Link to the paper : https://arxiv.org/abs/2102.12660
      """
      name = "Distributionally Robust Federated Averaging"
  
      def __init__(self, scenario, **kwargs):
          super(DistributionallyRobustFederatedAveragingLearning, self).__init__(scenario, **kwargs)
          if self.partners_count == 1:
              raise ValueError('Only one partner is provided. Please use the dedicated SinglePartnerLearning class')
          self.active_partners_count = scenario.active_partners_count
  
          self.lambda_vector = self.init_lambda()
          self.active_partners_list = list()
          self.update_active_partners_list()
  
          self.local_steps = scenario.gradient_updates_per_pass_count
          self.partners_training_data = {}
          self.partners_participation = self.initialize_participation_dict()
          self.lambda_learning_rate = 8e-3
  
          self.local_steps_index = 0
          self.local_steps_index_t = 0
          self.global_model_at_index_t = None
          self.model_weights_at_index_t = list()
          self.loss_for_model_at_index_t = np.zeros(self.partners_count)
  
          self.subset_u_partners = list()
          self.loss_vector_v = list()
  
      def fit_epoch(self):
  
          # Split the train dataset in mini-batches
          self.split_in_minibatches()
  
          # convert partners training data into tf Dataset, reference: fast_mpl
          for partner_id, partner in enumerate(self.partners_list):
              self.partners_training_data[partner.id] = list()
              for minibatch_index in range(self.minibatch_count):
                  # convert training data
                  data_train = tf.data.Dataset.from_tensor_slices((partner.minibatched_x_train[minibatch_index],
                                                                   partner.minibatched_y_train[minibatch_index]))
                  data_train = data_train.shuffle(len(partner.minibatched_x_train[minibatch_index]))
                  data_train = data_train.batch(partner.batch_size)
                  data_train = data_train.prefetch(1)
                  self.partners_training_data[partner.id].append(data_train)
  
          # Iterate over mini-batches and train
          for i in range(self.minibatch_count):
              self.minibatch_index = i
  
              self.local_steps_index = 0
              self.local_steps_index_t = np.random.randint(0, self.local_steps - 1)
  
              logger.info(
                  f"Active partner in this round "
                  f"{['#'+str(active_partner.id) for active_partner in self.active_partners_list]} "
                  f"according to lambda vector > {self.lambda_vector}")
              logger.info(f"Local step index t > {self.local_steps_index_t}")
  
              self.fit_minibatch()
  
              # update partner participations
              self.partners_participation[self.epoch_index][self.minibatch_index][[p.id for p
                                                                                   in self.active_partners_list]] = 1
  
              self.update_lambda()
              self.update_active_partners_list()
          self.log_partners_participation_rate()
  
          self.minibatch_index = 0
  
      def fit_minibatch(self):
          """Proceed to a collaborative round with a distributionally robust federated averaging approach"""
  
          # Starting model for each partner is the aggregated model from the previous mini-batch iteration
          logger.info(f"(drfa) Minibatch n°{self.minibatch_index} of epoch n°{self.epoch_index}, "
                      f"init each partner's models with a copy of the global model")
  
          for partner in self.partners_list:
              partner.model_weights = self.model_weights
  
          # Evaluate and store accuracy of mini-batch start model
          self.eval_and_log_model_val_perf()
  
          # Iterate over partners for training
          for partner_index, partner in enumerate(self.active_partners_list):
              partner_model = partner.build_model()
              # loop through each partner's minibatch
              minibatched_x_y = self.partners_training_data[partner.id][self.minibatch_index]
              for idx, batch_x_y in enumerate(minibatched_x_y):
                  with tf.GradientTape() as tape:
                      p_pred = partner_model(batch_x_y[0])
                      loss = partner_model.compiled_loss(batch_x_y[1], p_pred)
  
                  partner_model.optimizer.minimize(loss, partner_model.trainable_weights, tape=tape)
  
                  self.local_steps_index += 1
                  if self.local_steps_index == self.local_steps_index_t:
                      # save model weights for each partner at local step t
                      self.model_weights_at_index_t.append(partner.model_weights)
  
              partner.model_weights = partner_model.get_weights()
              self.local_steps_index = 0
  
          # aggregate final global model weights
          self.model_weights = self.aggregate_model_weights(self.active_partners_list)
  
          # build the model for each partner using weights gathered at index t
          for active_partner, weights_t in zip(self.active_partners_list, self.model_weights_at_index_t):
              active_partner.model_weights = weights_t
  
          # aggregate global model weights at index t
          self.global_model_at_index_t = self.aggregate_model_weights(self.active_partners_list)
  
          # sample a new subset of partners of size active_partners_count
          subset_index = random.sample(range(self.partners_count), self.active_partners_count)
          self.subset_u_partners = [self.partners_list[index] for index in subset_index]
          logger.info(
              f"Subset of partners chosen for lambda update "
              f"{['#'+ str(partner.id) for partner in self.subset_u_partners]}")
  
          # compute losses over a random batch using the global model at index t
          for partner, index in zip(self.subset_u_partners, subset_index):
              random_minibatch_index = np.random.randint(0, self.minibatch_count - 1)
              random_minibatch = self.partners_training_data[partner.id][random_minibatch_index]
              random_batch_index = np.random.randint(0, len(random_minibatch) - 1)
              random_batch = list(random_minibatch)[random_batch_index]
              partner_model = self.build_model_from_weights(self.global_model_at_index_t)
              loss = partner_model.compiled_loss(random_batch[1], partner_model(random_batch[0]))
              # compute (n/m)*loss and add it to the loss vector
              # n is the total number of partners, m is the number of active partners
              self.loss_for_model_at_index_t[index] = \
                  ((self.partners_count / self.active_partners_count) * np.mean(loss.numpy()))
  
      def init_lambda(self):
          """
          - initialize lambda vector according to each partner's dataset size
          - this is  a probability vector of size partners_count
          """
          return np.array(self.amounts_per_partner)
  
      def update_lambda(self):
          """
          The update rule for lambda is : lambda_vector(i) =
          Projection(lambda_vector(i-1) + (local_step_index_t * lambda_learning_rate * local_losses_at_index_t))
          """
          self.lambda_vector += (self.local_steps_index_t * self.lambda_learning_rate * self.loss_for_model_at_index_t)
          self.lambda_vector = project_onto_the_simplex(self.lambda_vector)
  
          # The projection can produce zero probabilities for certain partners which prevents them from
          # participating in the training. To avoid this, we assign 1e-3 to each probability smaller than this value.
          if any(self.lambda_vector < 1e-3):
              self.lambda_vector[self.lambda_vector < 1e-3] = 1e-3
              # normalize the probability vector
              self.lambda_vector = self.lambda_vector / np.sum(self.lambda_vector)
  
      def update_active_partners_list(self):
          """
          Update the active partners list according to lambda vector
          """
          active_partners_indices = (-self.lambda_vector).argsort()[:self.active_partners_count]
          self.active_partners_list = [self.partners_list[index] for index in active_partners_indices]
  
      def initialize_participation_dict(self):
          participation = {}
          for epoch_index in range(self.epoch_count):
              participation[epoch_index] = {}
              for minibatch_index in range(self.minibatch_count):
                  participation[epoch_index][minibatch_index] = np.zeros(self.partners_count)
          return participation
  
      def log_partners_participation_rate(self):
          epoch_participation_vector = np.zeros(self.partners_count)
          percentages = []
          for minibatch_index, vect in self.partners_participation[self.epoch_index].items():
              epoch_participation_vector += vect
              percentages = [str(np.round(p_v / self.minibatch_count, 2) * 100) + ' %'
                             for p_v in list(epoch_participation_vector)]
          logger.info(f"Partners {['#' + str(p.id) for p in self.partners_list]} "
                      f"have the following participation rates, respectively : "
                      f"{percentages} "
                      f"at the end of Epoch > {self.epoch_index}")
  
          final_participation_vector = np.zeros(self.partners_count)
          if self.epoch_index == self.epoch_count - 1:
              for epoch_index in range(self.epoch_count):
                  for minibatch_index, vect in self.partners_participation[epoch_index].items():
                      final_participation_vector += vect
                      percentages = [str(np.round(f_p_v / (self.minibatch_count * self.epoch_count), 2) * 100) + '%'
                                     for f_p_v in list(final_participation_vector)]
              logger.info(f"Partners {['#' + str(p.id) for p in self.partners_list]} "
                          f"have the following participation rates : "
                          f"{percentages} "
                          f"during the training")
  
      @staticmethod
      def aggregate_model_weights(partners_list):
          """ This method is identical to the one in the aggregator class with few modifications.
           I couldn't use the original aggregator method since it operates on the entire list of partners and
           DRFA requires model aggregation over a subset of partners list only
          """
          aggregation_weights = np.ones(len(partners_list), dtype='float32')
          weights_per_layer = list(zip(*[partner.model_weights for partner in partners_list]))
          new_weights = list()
  
          for weights_for_layer in weights_per_layer:
              avg_weights_for_layer = np.average(
                  np.array(weights_for_layer), axis=0, weights=aggregation_weights
              )
              new_weights.append(avg_weights_for_layer)
  
          return new_weights
  
  
  class SequentialLearning(MultiPartnerLearning):  # seq-pure
      name = 'Sequential learning'
  
      def __init__(self, scenario, **kwargs):
          super(SequentialLearning, self).__init__(scenario, **kwargs)
          if self.partners_count == 1:
              raise ValueError('Only one partner is provided. Please use the dedicated SinglePartnerLearning class')
  
      def fit_epoch(self):
          # Clear Keras' old models
          clear_session()
  
          # Split the train dataset in mini-batches
          self.split_in_minibatches()
  
          # Iterate over mini-batches and train
          for i in range(self.minibatch_count):
              self.minibatch_index = i
              logger.info(f"(seq-pure) Minibatch n°{self.minibatch_index} of epoch n°{self.epoch_index}")
              self.fit_minibatch()
  
      def fit_minibatch(self):
          """Proceed to a collaborative round with a sequential averaging approach"""
  
          logger.debug("Start new seq collaborative round ...")
  
          model_for_round = self.build_model()
  
          # Evaluate and store accuracy of mini-batch start model
          self.eval_and_log_model_val_perf()
          # Iterate over partners for training each individual model
          shuffled_indexes = np.random.permutation(self.partners_count)
          logger.debug(f"(seq) Shuffled order for this seqavg collaborative round: {shuffled_indexes}")
          for idx, partner_index in enumerate(shuffled_indexes):
              partner = self.partners_list[partner_index]
  
              # Train on partner local data set
              if self.val_set == 'global':
                  history = model_for_round.fit(partner.minibatched_x_train[self.minibatch_index],
                                                partner.minibatched_y_train[self.minibatch_index],
                                                batch_size=partner.batch_size,
                                                verbose=0,
                                                validation_data=self.val_data)
              elif self.val_set == 'local':
                  history = model_for_round.fit(partner.minibatched_x_train[self.minibatch_index],
                                                partner.minibatched_y_train[self.minibatch_index],
                                                batch_size=partner.batch_size,
                                                verbose=0,
                                                validation_data=(partner.x_val, partner.y_val))
              else:
                  raise ValueError("validation set should be 'local' or 'global', not {self.val_set}")
  
              # Log results
              self.log_partner_perf(partner.id, idx, history.history)
  
              # Save the partner's model in the models' list
              partner.model_weights = model_for_round.get_weights()
              self.model_weights = model_for_round.get_weights()
  
          logger.debug("End of seq collaborative round.")
  
  
  class SequentialWithFinalAggLearning(SequentialLearning):
      name = 'Sequential learning with final aggregation'
  
      def __init__(self, scenario, **kwargs):
          super(SequentialWithFinalAggLearning, self).__init__(scenario, **kwargs)
          if self.partners_count == 1:
              raise ValueError('Only one partner is provided. Please use the dedicated SinglePartnerLearning class')
  
      def fit_epoch(self):
          # Clear Keras' old models
          clear_session()
  
          # Split the train dataset in mini-batches
          self.split_in_minibatches()
  
          # Iterate over mini-batches and train
          for i in range(self.minibatch_count):
              logger.info(f"(seq-final-agg) Minibatch n°{self.minibatch_index} of epoch n°{self.epoch_index}, "
                          f"init model with a copy of the global model")
              self.minibatch_index = i
              self.fit_minibatch()
  
          # At the end of each epoch, aggregate the models
          self.model_weights = self.aggregator.aggregate_model_weights()
  
  
  class SequentialAverageLearning(SequentialLearning):
      name = 'Sequential averaged learning'
  
      def __init__(self, scenario, **kwargs):
          super(SequentialAverageLearning, self).__init__(scenario, **kwargs)
          if self.partners_count == 1:
              raise ValueError('Only one partner is provided. Please use the dedicated SinglePartnerLearning class')
  
      def fit_epoch(self):
          # Clear Keras' old models
          clear_session()
  
          # Split the train dataset in mini-batches
          self.split_in_minibatches()
  
          # Iterate over mini-batches and train
          for i in range(self.minibatch_count):
              logger.info(f"(seqavg) Minibatch n°{self.minibatch_index} of epoch n°{self.epoch_index}, "
                          f"init model with a copy of the global model")
              self.minibatch_index = i
              self.fit_minibatch()
  
              # At the end of each minibatch, aggregate the models
              self.model_weights = self.aggregator.aggregate_model_weights()
  
  
  class FedAvgSmodel(FederatedAverageLearning):
      name = 'Federated learning with label flipping'
  
      def __init__(self, scenario, pretrain_epochs=0, epsilon=0.5, **kwargs):
          super(FedAvgSmodel, self).__init__(scenario, **kwargs)
          self.pretrain_epochs = pretrain_epochs
          self.epsilon = epsilon
          if pretrain_epochs > 0:
              self.pretrain_mpl = FederatedAverageLearning(scenario=scenario,
                                                           epoch_count=self.pretrain_epochs,
                                                           is_save_data=False)
  
      def fit(self):
          if self.pretrain_epochs > 0:
              logger.info('Start pre-train...')
              self.pretrain_mpl.fit()
              pretrain_model = self.pretrain_mpl.build_model()
              for p in self.partners_list:
                  confusion = confusion_matrix(np.argmax(p.y_train, axis=1),
                                               np.argmax(pretrain_model.predict(p.x_train), axis=1),
                                               normalize='pred')
                  p.noise_layer_weights = [np.log(confusion.T + 1e-8)]
              self.model_weights[:-1] = self.pretrain_mpl.model_weights[:-1]
          else:
              for p in self.partners_list:
                  confusion = np.identity(10) * (1 - self.epsilon) + (self.epsilon / 10)
                  p.noise_layer_weights = [np.log(confusion + 1e-8)]
          super(FedAvgSmodel, self).fit()
  
      def fit_minibatch(self):
          """Proceed to a collaborative round with a S-Model federated averaging approach"""
  
          logger.debug("Start new S-Model collaborative round ...")
  
          # Starting model for each partner is the aggregated model from the previous mini-batch iteration
          logger.info(f"(S-Model) Minibatch n°{self.minibatch_index} of epoch n°{self.epoch_index}, "
                      f"init each partner's models with a copy of the global model")
  
          for partner in self.partners_list:
              partner.model_weights = self.model_weights
  
          # Evaluate and store accuracy of mini-batch start model
          self.eval_and_log_model_val_perf()
  
          # Iterate over partners for training each individual model
          for partner_index, partner in enumerate(self.partners_list):
              # Reference the partner's model
              partner_model = partner.build_model()
              x_batch = partner.minibatched_x_train[self.minibatch_index]
              y_batch = partner.minibatched_y_train[self.minibatch_index]
  
              model_input = Input(shape=self.dataset.input_shape)
              x = partner_model(model_input)
              outputs = NoiseAdaptationChannel(weights=partner.noise_layer_weights, name='s-model')(x)
              full_model = Model(inputs=model_input, outputs=outputs, name=f"full_model_partner_{partner_index}")
  
              full_model.compile(
                  loss=partner_model.loss,
                  optimizer=partner_model.optimizer,
                  metrics='accuracy',
              )
  
              # Train on partner local data set
              history = full_model.fit(x_batch,
                                       y_batch,
                                       batch_size=partner.batch_size,
                                       verbose=0,
                                       validation_data=self.val_data)
  
              # Log results of the round
              self.log_partner_perf(partner.id, partner_index, history.history)
  
              # Update the partner's model in the models' list
              partner.noise_layer_weights = full_model.get_layer('s-model').get_weights()
              partner.model_weights = partner_model.get_weights()
  
          logger.debug("End of S-Model collaborative round.")
  
  
  class FederatedGradients(MultiPartnerLearning):
      def __init__(self, scenario, **kwargs):
          super(FederatedGradients, self).__init__(scenario, **kwargs)
          if self.partners_count == 1:
              raise ValueError('Only one partner is provided. Please use the dedicated SinglePartnerLearning class')
          self.model = self.build_model()
  
      def fit_epoch(self):
          # Split the train dataset in mini-batches
          self.split_in_minibatches()
          # Iterate over mini-batches and train
          for i in range(self.minibatch_count):
              self.minibatch_index = i
              self.fit_minibatch()
  
          self.minibatch_index = 0
  
      def fit_minibatch(self):
          """Proceed to a collaborative round with a federated averaging approach"""
  
          logger.debug("Start new gradients fusion collaborative round ...")
  
          # Starting model for each partner is the aggregated model from the previous mini-batch iteration
          logger.info(f"(gradient fusion) Minibatch n°{self.minibatch_index} of epoch n°{self.epoch_index}, "
                      f"init each partner's models with a copy of the global model")
  
          for partner in self.partners_list:
              # Evaluate and store accuracy of mini-batch start model
              partner.model_weights = self.model_weights
          self.eval_and_log_model_val_perf()
  
          # Iterate over partners for training each individual model
          for partner_index, partner in enumerate(self.partners_list):
              with tf.GradientTape() as tape:
                  loss = self.model.loss(partner.minibatched_y_train[self.minibatch_index],
                                         self.model(partner.minibatched_x_train[self.minibatch_index]))
              partner.grads = tape.gradient(loss, self.model.trainable_weights)
  
          global_grad = self.aggregator.aggregate_gradients()
          self.model.optimizer.apply_gradients(zip(global_grad, self.model.trainable_weights))
          self.model_weights = self.model.get_weights()
  
          for partner_index, partner in enumerate(self.partners_list):
              val_history = self.model.evaluate(self.val_data[0], self.val_data[1], verbose=False)
              history = self.model.evaluate(partner.minibatched_x_train[self.minibatch_index],
                                            partner.minibatched_y_train[self.minibatch_index], verbose=False)
              history = {
                  "loss": [history[0]],
                  'accuracy': [history[1]],
                  'val_loss': [val_history[0]],
                  'val_accuracy': [val_history[1]]
              }
  
              # Log results of the round
              self.log_partner_perf(partner.id, partner_index, history)
  
          logger.debug("End of grads-fusion collaborative round.")
  
  
  class EnsemblePredictions(MultiPartnerLearning):
      """
      Ensemble (average) prediction of several input models
      This approach can only be used with the EnsemblePredictionsModel
      """
  
      def __init__(self, scenario, **kwargs):
          super(EnsemblePredictions, self).__init__(scenario, **kwargs)
  
          # First, if only one partner, fall back to dedicated single partner function
          if self.partners_count == 1:
              raise ValueError('Only one partner is provided. Please use the dedicated SinglePartnerLearning class')
  
          partner_model_list = [self.dataset.generate_new_model() for _ in range(self.partners_count)]
          self.model = EnsemblePredictionsModel(partner_model_list)
  
          for partner in self.partners_list:
              partner.model_weights = deepcopy(self.model_weights)
              print(id(partner.model_weights))
  
          logger.info("Init EnsemblePredictionsModel model")
  
      def build_model(self):
          partner_model_list = [partner.build_model() for partner in self.partners_list]
          return EnsemblePredictionsModel(partner_model_list)
  
      def fit_epoch(self):
          # Clear Keras' old models
          clear_session()
  
          self.eval_and_log_model_val_perf()
  
          for partner_index, partner in enumerate(self.partners_list):
  
              partner_model = partner.build_model()
  
              # Train on partner local data set
              history = partner_model.fit(partner.x_train,
                                          partner.y_train,
                                          batch_size=partner.batch_size,
                                          verbose=0,
                                          validation_data=self.val_data)
  
              # Log results of the round
              self.log_partner_perf(partner.id, partner_index, history.history)
  
              # Update the partner's model in the models' list
              partner.model_weights = partner_model.get_weights()
  
      def fit_minibatch(self):
          pass
  
apis: numpy.sum numpy.zeros numpy.random.permutation tensorflow.keras.backend.clear_session tensorflow.keras.Model numpy.argmax tensorflow.GradientTape numpy.log tensorflow.keras.callbacks.EarlyStopping numpy.identity numpy.array tensorflow.data.Dataset.from_tensor_slices numpy.random.randint numpy.round tensorflow.keras.Input