求助：如何学习预训练神经网络权重的加权平均？

你的核心需求是仅学习加权系数，对多个同架构预训练模型的权重做加权平均，得到新模型并提升泛化能力，原代码存在两个关键问题：call方法内修改原模型权重的副作用，以及权重计算时的维度错误。以下是修正后的实现方案：

核心思路 #

1. 固定所有预训练模型的权重，仅将集成加权系数设为可训练变量
2. 使用softmax约束加权系数，保证其非负且和为1，符合加权平均逻辑
3. 用克隆的模板模型执行前向传播，避免修改原模型权重
4. 用tf.stack+tf.einsum处理任意维度的参数加权，解决形状不兼容问题

完整代码实现 #

import tensorflow as tf
from tensorflow.keras import layers, initializers, models

class Ensemble(layers.Layer):
    def __init__(self, model_list):
        super().__init__()
        self.num_models = len(model_list)
        
        # 保存所有预训练模型的权重，转为非可训练变量（避免被更新）
        self.model_weights = []
        for model in model_list:
            fixed_weights = [tf.Variable(w, trainable=False) for w in model.get_weights()]
            self.model_weights.append(fixed_weights)
        
        # 克隆一个模板模型，用于每次前向传播（不修改原模型）
        self.template_model = models.clone_model(model_list[0])
        # 初始化模板模型权重（后续会被加权权重覆盖）
        self.template_model.set_weights(model_list[0].get_weights())
        
        # 定义可训练的集成加权系数，初始化为均匀分布
        self.w = self.add_weight(
            name='ensemble_weights',
            shape=(self.num_models,),
            initializer=initializers.Constant(1.0 / self.num_models),
            trainable=True
        )

    def call(self, inputs):
        # 对加权系数做softmax，确保非负且和为1
        normalized_w = tf.nn.softmax(self.w)
        
        # 计算每个参数的加权平均
        averaged_weights = []
        # 遍历所有模型的对应位置参数（如所有模型的第1层kernel）
        for params in zip(*self.model_weights):
            # 将N个模型的同位置参数堆叠为(N, ...)形状的张量
            params_stack = tf.stack(params, axis=0)
            # 加权平均，自动适配任意维度的参数
            averaged_param = tf.einsum('n,...->...', normalized_w, params_stack)
            averaged_weights.append(averaged_param)
        
        # 设置模板模型的权重为加权后的结果，执行前向传播
        self.template_model.set_weights(averaged_weights)
        return self.template_model(inputs)

使用示例 #

假设你已拥有3个同架构的预训练模型model1、model2、model3：

# 构建集成模型
inputs = layers.Input(shape=(你的输入形状,))
ensemble_output = Ensemble([model1, model2, model3])(inputs)
ensemble_model = models.Model(inputs=inputs, outputs=ensemble_output)

# 编译并训练（仅更新加权系数w）
ensemble_model.compile(optimizer='adam', loss='categorical_crossentropy', metrics=['accuracy'])
ensemble_model.fit(x_train, y_train, epochs=10, validation_data=(x_val, y_val))

原代码问题分析 #

1. call方法内修改原模型权重：self.model.set_weights()会直接修改原模型的权重，导致原模型被污染，且Keras的call方法要求无副作用，这种操作会破坏训练计算图的正确性。
2. 权重计算的维度错误：嵌套的zip(*weightsTuple)会错误拆解参数的维度（如将形状为(3,4,64)的卷积核拆分为3个(4,64)的张量），导致加权后的参数形状与模型要求不匹配，触发ValueError。
3. 多余的偏置b：加权平均已包含原模型的偏置参数，额外添加全局偏置会破坏原模型的结构逻辑，无必要。

可选优化：动态层权重分配 #

如果模型规模较大，每次调用set_weights存在性能开销，可以直接构建层列表并动态分配加权权重，避免克隆模型：

class Ensemble(layers.Layer):
    def __init__(self, model_list):
        super().__init__()
        self.num_models = len(model_list)
        self.model_weights = []
        for model in model_list:
            fixed_weights = [tf.Variable(w, trainable=False) for w in model.get_weights()]
            self.model_weights.append(fixed_weights)
        
        self.w = self.add_weight(
            name='ensemble_weights',
            shape=(self.num_models,),
            initializer=initializers.Constant(1.0 / self.num_models),
            trainable=True
        )
        
        # 构建与原模型一致的层列表
        self.layers_list = []
        original_layers = model_list[0].layers
        for layer in original_layers:
            if isinstance(layer, layers.Dense):
                new_layer = layers.Dense(units=layer.units, activation=layer.activation, use_bias=layer.use_bias)
                new_layer.build((None,) + layer.input_shape[1:])
                self.layers_list.append(new_layer)
            elif isinstance(layer, layers.Conv2D):
                new_layer = layers.Conv2D(filters=layer.filters, kernel_size=layer.kernel_size, 
                                         strides=layer.strides, padding=layer.padding, 
                                         activation=layer.activation, use_bias=layer.use_bias)
                new_layer.build((None,) + layer.input_shape[1:])
                self.layers_list.append(new_layer)
            elif isinstance(layer, layers.Flatten):
                self.layers_list.append(layers.Flatten())
            else:
                self.layers_list.append(layer.__class__.from_config(layer.get_config()))

    def call(self, inputs):
        normalized_w = tf.nn.softmax(self.w)
        averaged_weights = []
        for params in zip(*self.model_weights):
            params_stack = tf.stack(params, axis=0)
            averaged_param = tf.einsum('n,...->...', normalized_w, params_stack)
            averaged_weights.append(averaged_param)
        
        # 动态分配权重并执行前向传播
        idx = 0
        x = inputs
        for layer in self.layers_list:
            if isinstance(layer, (layers.Dense, layers.Conv2D)):
                layer.kernel.assign(averaged_weights[idx])
                if layer.use_bias:
                    layer.bias.assign(averaged_weights[idx+1])
                    idx +=2
                else:
                    idx +=1
                x = layer(x)
            else:
                x = layer(x)
        return x

内容的提问来源于stack exchange，提问作者Ghull