Learn a Weighted Average of Pre-Trained Neural Network Weights

Question

I have multiple pre-trained neural networks with the same architecture and different weights. I want to take a weighted average of the weights in these networks to make one network of the same size (in an attempt to improve generalizability without sacrificing size).

To be clear: I only want to learn the weights for the average: NOT the weights inside the pre-trained networks.

This is what I have so far:

class Ensemble(layers.Layer):
    def __init__(self, modelWeights, model):
        super().__init__()

        self.modelWeights = modelWeights
        self.model = model

        self.w = self.add_weight(f'Weights', shape = (len(modelWeights),), initializer = initializers.Constant(1 / len(modelWeights)))
        self.b = self.add_weight('Bias', shape = (1,), initializer = 'zeros')

    def call(self, inputs):
        newWeights = []
        for weightsTuple in zip(*self.modelWeights):
            temp = []
            for weights in zip(*weightsTuple):
                weights = tf.convert_to_tensor(weights)
                temp += [tf.tensordot(weights, self.w, axes = [[0], [0]]) + self.b]

            newWeights += [temp]
    
        self.model.set_weights(newWeights)
        return self.model(inputs)

modelWeights is a list of model.get_weights()

Besides the error I'm currently getting (ValueError: Layer model weight shape (3, 4, 64) is not compatible with provided weight shape ()) I don't think keras is going to let me do self.model.set_weights(newWeights) inside the call function.

Does anyone have a better way to do this?

Thanks in advance

Answer 1

I hope I get your idea, otherwise correct me.
To average out the weights of multiple trained models you can do the following (example with 3 models):

models_weights = [model1.get_weights(), model2.get_weights(), model3.get_weights()]
models_scalars = [0.3, 0.2, 0.5]  # scalar weighting for model1, model2..

new_model_weights = [[layer_weight * scalar for layer_weight in model_weights] for model_weights, scalar in zip(models_weights, model_scalars)]
new_model_weights_mean = [np.sum(np.array(layer_weights),axis=0) for layer_weights in zip(*new_model_weights)]
# create new_model here..
new_model.save_weights(new_model_weights_mean)

Edit: The only other approach I can think of is to freeze all models and connect their output to a Dense(1, use_bias=False) node, with a kernel constrain that all weights has to be between 0 and 1 and sum up to 1. This way a model could potentially learn how much say has a network in the final decision. But this will not merge the network weights together and only weights the outputs.