如何在Keras层实现双输入权重调制网络（复现AI4Animation ExpertModel）

问题根源分析 #

你当前自定义层的核心错误是在call方法内创建tf.Variable：每次调用层都会生成新的可训练变量，这些变量无法被模型的梯度追踪系统正确关联，直接导致反向传播失效。此外，这种方式也无法让门控网络的输出（weight_blend）与主网络的权重建立可微分的关联。

正确实现思路 #

1. 预初始化专家权重：主网络的所有专家（Expert）的权重和偏置，需要在层的__init__方法中预先定义为可训练变量，而非在call中动态创建。
2. 张量级权重混合：在call方法中，用门控网络输出的weight_blend系数，对预定义的专家权重/偏置进行加权求和，得到当前批次使用的混合权重——这一步是纯张量运算，可被TensorFlow自动微分系统追踪，保证反向传播有效。
3. 合并双输入模型：将主输入、门控输入作为模型的两个独立输入，门控网络输出的weight_blend直接传入主网络的自定义层，形成端到端的可训练模型。

具体代码实现 #

1. 自定义权重混合全连接层 #

import tensorflow as tf

class WeightBlendedDense(tf.keras.layers.Layer):
    def __init__(self, num_experts, units, activation=None, **kwargs):
        super().__init__(**kwargs)
        self.num_experts = num_experts
        self.units = units
        self.activation = tf.keras.activations.get(activation)
        
        # 预初始化所有专家的权重和偏置：形状为 [num_experts, input_dim, units] / [num_experts, units]
        # 注意：input_dim需要在build阶段确定
        self.expert_weights = None
        self.expert_biases = None

    def build(self, input_shape):
        # input_shape是一个列表：[主输入形状, weight_blend形状]
        main_input_shape = input_shape[0]
        input_dim = main_input_shape[-1]
        
        # 初始化专家权重：num_experts个 [input_dim, units] 的权重矩阵
        self.expert_weights = self.add_weight(
            shape=(self.num_experts, input_dim, self.units),
            initializer="glorot_uniform",
            trainable=True,
            name="expert_weights"
        )
        # 初始化专家偏置：num_experts个 [units] 的偏置向量
        self.expert_biases = self.add_weight(
            shape=(self.num_experts, self.units),
            initializer="zeros",
            trainable=True,
            name="expert_biases"
        )
        super().build(input_shape)

    def call(self, inputs):
        x, weight_blend = inputs
        
        # weight_blend形状：[batch_size, num_experts]，需保证已做归一化（如softmax）
        # 对专家权重进行加权混合：
        # weight_blend扩展维度为 [batch_size, num_experts, 1, 1]
        # 与expert_weights（[num_experts, input_dim, units]）相乘后求和，得到 [batch_size, input_dim, units]
        blended_weights = tf.reduce_sum(
            weight_blend[:, :, tf.newaxis, tf.newaxis] * self.expert_weights,
            axis=1
        )
        # 混合偏置：[batch_size, num_experts, 1] * [num_experts, units] → [batch_size, units]
        blended_biases = tf.reduce_sum(
            weight_blend[:, :, tf.newaxis] * self.expert_biases,
            axis=1
        )
        
        # 执行全连接运算：x形状 [batch_size, input_dim] → 转置为 [batch_size, 1, input_dim]
        # 与blended_weights [batch_size, input_dim, units] 相乘 → [batch_size, 1, units] →  squeeze为 [batch_size, units]
        x = tf.matmul(x[:, tf.newaxis, :], blended_weights)[:, 0, :] + blended_biases
        
        if self.activation is not None:
            x = self.activation(x)
        return x

2. 构建完整的ExpertModel #

def build_expert_model(num_experts, main_input_dim, gating_input_dim, hidden_units):
    # 定义两个独立输入
    main_input = tf.keras.Input(shape=(main_input_dim,), name="main_input")
    gating_input = tf.keras.Input(shape=(gating_input_dim,), name="gating_input")
    
    # 门控网络：输出num_experts维的混合系数，用softmax归一化
    gating_network = tf.keras.Sequential([
        tf.keras.layers.Dense(64, activation="relu"),
        tf.keras.layers.Dense(num_experts, activation="softmax")
    ], name="gating_network")
    weight_blend = gating_network(gating_input)
    
    # 主网络：使用自定义权重混合层
    main_network = tf.keras.Sequential([
        WeightBlendedDense(num_experts, hidden_units, activation="relu"),
        WeightBlendedDense(num_experts, hidden_units, activation="relu"),
        WeightBlendedDense(num_experts, 1)  # 示例输出维度，根据任务调整
    ], name="main_network")
    output = main_network([main_input, weight_blend])
    
    # 构建端到端模型
    model = tf.keras.Model(inputs=[main_input, gating_input], outputs=output)
    return model

3. 训练示例 #

# 初始化模型
model = build_expert_model(
    num_experts=5,
    main_input_dim=10,
    gating_input_dim=8,
    hidden_units=32
)

# 编译模型
model.compile(optimizer="adam", loss="mse")

# 模拟数据
batch_size = 32
main_data = tf.random.normal((1000, 10))
gating_data = tf.random.normal((1000, 8))
labels = tf.random.normal((1000, 1))

# 训练：只需传入两个输入和标签，梯度会自动回传到主输入和门控输入
model.fit(
    x={"main_input": main_data, "gating_input": gating_data},
    y=labels,
    batch_size=batch_size,
    epochs=10
)

关键说明 #

• 所有专家权重均在__init__/build阶段初始化，属于模型的可训练参数，会被梯度下降正确更新。
• 权重混合过程是纯张量运算，TensorFlow的自动微分系统可以完整追踪从输出到主输入、门控输入的梯度路径，解决了你之前反向传播失效的问题。
• 模型保留了两个独立输入，无需使用Concatenate层，符合你对输入分离的需求。

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