如何在自定义TensorFlow Layer中获取当前学习率、epoch或step？

以下是几种无需在Model中遍历调用的实现方式，适配你的需求：

方案1：给自定义Layer注入优化器引用 #

通过在Model的compile方法中，将优化器引用传递给所有自定义Layer，让Layer直接访问优化器的lr（学习率）和iterations（全局步数）属性。

自定义Layer修改 #

class my_layer(tf.keras.layers.Layer):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.optimizer = None  # 初始化优化器引用
    
    # 保留原有constructor、build方法
    def function_to_get_lr(self, lr):
        # 这里处理学习率相关逻辑
        print(f"当前学习率：{lr}")
    
    def call(self, inputs, training=False):
        if training and self.optimizer is not None:
            # 获取全局步数
            current_step = tf.keras.backend.get_value(self.optimizer.iterations)
            # 区分固定学习率和学习率调度器的情况
            if isinstance(self.optimizer.lr, tf.keras.optimizers.schedules.LearningRateSchedule):
                current_lr = self.optimizer.lr(current_step).numpy()
            else:
                current_lr = tf.keras.backend.get_value(self.optimizer.lr)
            self.function_to_get_lr(current_lr)
        # 原有call逻辑
        return inputs

自定义Model修改 #

class my_model(tf.keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.layer1 = my_layer(name="my_layer_1")
        self.layer2 = my_layer(name="my_layer_2")
        # 其他层定义
    
    def compile(self, optimizer, **kwargs):
        super().compile(optimizer=optimizer, **kwargs)
        # 给所有my_layer实例注入优化器引用
        for layer in self.layers:
            if isinstance(layer, my_layer):
                layer.optimizer = self.optimizer
    
    def call(self, inputs, training=False):
        x = self.layer1(inputs, training=training)
        x = self.layer2(x, training=training)
        # 其他前向逻辑
        return x

方案2：通过自定义训练循环传递步数/学习率 #

如果使用自定义训练循环而非Keras内置的fit方法，可以直接将当前步数或学习率作为参数传入Layer的call方法，无需依赖优化器引用。

自定义Layer修改 #

class my_layer(tf.keras.layers.Layer):
    # 保留原有constructor、build方法
    def function_to_get_lr(self, lr):
        # 处理学习率逻辑
        print(f"当前学习率：{lr}")
    
    def call(self, inputs, training=False, current_step=None):
        if training and current_step is not None:
            # 根据步数计算当前学习率（适配学习率调度器）
            current_lr = self.optimizer.lr(current_step).numpy()
            self.function_to_get_lr(current_lr)
        # 原有call逻辑
        return inputs

自定义训练循环示例 #

# 初始化模型、优化器、损失函数
model = my_model()
optimizer = tf.keras.optimizers.Adam(learning_rate=tf.keras.optimizers.schedules.ExponentialDecay(
    initial_learning_rate=0.01, decay_steps=1000, decay_rate=0.9
))
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 全局步数变量
global_step = tf.Variable(0, trainable=False, dtype=tf.int64)

# 训练循环
epochs = 10
for epoch in range(epochs):
    print(f"Epoch {epoch+1}/{epochs}")
    for x_batch, y_batch in train_dataset:
        with tf.GradientTape() as tape:
            # 传入当前步数到Layer
            y_pred = model(x_batch, training=True, current_step=global_step)
            loss = loss_fn(y_batch, y_pred)
        
        # 计算梯度并更新参数
        grads = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(grads, model.trainable_variables))
        
        # 更新全局步数
        global_step.assign_add(1)

方案3：共享全局步数变量 #

在Model中创建一个全局步数变量，初始化自定义Layer时传入该变量，Layer可直接通过变量获取当前步数，再计算学习率。

自定义Layer修改 #

class my_layer(tf.keras.layers.Layer):
    def __init__(self, global_step, **kwargs):
        super().__init__(**kwargs)
        self.global_step = global_step
        self.optimizer = None  # 若需计算学习率则注入
    
    def function_to_get_lr(self, lr):
        # 处理学习率逻辑
        print(f"当前步数：{tf.keras.backend.get_value(self.global_step)}，学习率：{lr}")
    
    def call(self, inputs, training=False):
        if training and self.optimizer is not None:
            current_step = tf.keras.backend.get_value(self.global_step)
            current_lr = self.optimizer.lr(current_step).numpy()
            self.function_to_get_lr(current_lr)
        # 原有call逻辑
        return inputs

自定义Model修改 #

class my_model(tf.keras.Model):
    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        self.global_step = tf.Variable(0, trainable=False, dtype=tf.int64)
        # 初始化Layer时传入全局步数变量
        self.layer1 = my_layer(global_step=self.global_step, name="my_layer_1")
        self.layer2 = my_layer(global_step=self.global_step, name="my_layer_2")
    
    def compile(self, optimizer, **kwargs):
        super().compile(optimizer=optimizer, **kwargs)
        # 注入优化器引用
        for layer in self.layers:
            if isinstance(layer, my_layer):
                layer.optimizer = self.optimizer
    
    def train_step(self, data):
        # 重写train_step方法更新全局步数（适配Keras fit方法）
        x, y = data
        with tf.GradientTape() as tape:
            y_pred = self(x, training=True)
            loss = self.compiled_loss(y, y_pred)
        
        grads = tape.gradient(loss, self.trainable_variables)
        self.optimizer.apply_gradients(zip(grads, self.trainable_variables))
        
        # 更新全局步数
        self.global_step.assign_add(1)
        
        # 更新指标
        self.compiled_metrics.update_state(y, y_pred)
        return {m.name: m.result() for m in self.metrics}

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