How to compute the mean of weights of multiple models?

Question

Hi i'm a student and i'm working on a Federated Learning problem, but before doing that with the proper tools like OpenFL or Flower, I started a little experiment to try in local to train using this technique.

I managed to train multiple models using IID data, now I'm struggling with the local_update() function that should collect the models and then i need to take all the weights of these models and compute their mean. I read some documentation of Keras and Tensorflow that I'm using for my work, and i found some functions but i can't get it to work properly.

Currently this is my local_update() that's not working

def local_update(self, models):

       
       weights = []
       #Take the weights of the models and compute the mean then return the weights to an updated model

       for model in models:
           for layer in model.layers:
               weights = layer.get_weights()
       
       #Compute the mean of weights  
       weights = np.mean(weights, axis=0)


       for layer in self.model.layers:
           self.model.set_weights(weights)     


       return self.model

In TensorFlow/Keras there are many way to do this but what is the best and simplest one?

Thank you in advance for the help!