如何在同一Keras层中为不同权重集设置不同学习率？

首先修正你的自定义层代码——原代码在__init__中直接使用input_shape和output_shape会报错，因为这些参数在层初始化阶段还未确定，正确做法是将权重定义放在build方法中：

import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers, optimizers

class MyCustomLayer(layers.Layer):
    def __init__(self):
        super(MyCustomLayer, self).__init__()

    def build(self, input_shape):
        # 输入形状确定后再创建权重
        self.set_1_weights = self.add_weight(
            shape=(input_shape[1], 10), 
            name='set_1_weights',
            trainable=True
        )
        # 这里假设输出维度与输入一致，可根据需求自定义
        self.set_2_weights = self.add_weight(
            shape=(10, input_shape[1]), 
            name='set_2_weights',
            trainable=True
        )
        super().build(input_shape)

    def call(self, inputs):
        x = tf.matmul(inputs, self.set_1_weights)
        output = tf.matmul(x, self.set_2_weights)
        return output

你可以通过自定义训练逻辑实现同一层内不同权重使用不同学习率的需求，以下是两种可行方式：

方法1：自定义训练循环分离权重更新 #

直接编写训练循环，分别计算两组权重的梯度并使用对应优化器更新：

1. 初始化两个不同学习率的Adam优化器：

opt_set1 = optimizers.Adam(learning_rate=0.01)
opt_set2 = optimizers.Adam(learning_rate=0.001)

2. 构建模型并定义损失函数：

input_layer = layers.Input(shape=(20,))
custom_layer = MyCustomLayer()
output_layer = custom_layer(input_layer)
model = keras.Model(inputs=input_layer, outputs=output_layer)

loss_fn = keras.losses.MeanSquaredError()

3. 编写训练循环：

# 示例训练数据
x_train = tf.random.normal((1000, 20))
y_train = tf.random.normal((1000, 20))

epochs = 10
batch_size = 32
steps_per_epoch = len(x_train) // batch_size

for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    for step in range(steps_per_epoch):
        start = step * batch_size
        end = start + batch_size
        x_batch = x_train[start:end]
        y_batch = y_train[start:end]

        # 记录梯度
        with tf.GradientTape(persistent=True) as tape:
            y_pred = model(x_batch, training=True)
            loss = loss_fn(y_batch, y_pred)
        
        # 分别获取两组权重的梯度
        grads_set1 = tape.gradient(loss, custom_layer.set_1_weights)
        grads_set2 = tape.gradient(loss, custom_layer.set_2_weights)
        
        # 用对应优化器更新权重
        opt_set1.apply_gradients([(grads_set1, custom_layer.set_1_weights)])
        opt_set2.apply_gradients([(grads_set2, custom_layer.set_2_weights)])
        
        if step % 10 == 0:
            print(f"Step {step}, Loss: {loss.numpy():.4f}")

方法2：自定义Model类适配model.fit接口 #

如果想保留Keras的model.fit便捷性，可以通过重写train_step方法实现：

class CustomModel(keras.Model):
    def __init__(self, custom_layer):
        super().__init__()
        self.custom_layer = custom_layer
        self.opt_set1 = optimizers.Adam(learning_rate=0.01)
        self.opt_set2 = optimizers.Adam(learning_rate=0.001)
        self.loss_fn = keras.losses.MeanSquaredError()

    def train_step(self, data):
        x, y = data
        with tf.GradientTape(persistent=True) as tape:
            y_pred = self.custom_layer(x, training=True)
            loss = self.loss_fn(y, y_pred)
        
        # 分别更新两组权重
        grads_set1 = tape.gradient(loss, self.custom_layer.set_1_weights)
        grads_set2 = tape.gradient(loss, self.custom_layer.set_2_weights)
        self.opt_set1.apply_gradients([(grads_set1, self.custom_layer.set_1_weights)])
        self.opt_set2.apply_gradients([(grads_set2, self.custom_layer.set_2_weights)])
        
        return {"loss": loss}

# 实例化并训练
custom_model = CustomModel(MyCustomLayer())
custom_model.compile()  # 无需传入优化器，已自定义训练逻辑
custom_model.fit(x_train, y_train, epochs=10, batch_size=32)

核心原理 #

通过tf.GradientTape单独记录损失对两组权重的梯度，再分别调用对应优化器的apply_gradients方法更新权重，从而实现同一层内不同权重的差异化学习率配置。

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