TensorFlow中带预热与余弦衰减的自定义学习率调度报错求助

Let’s break down why you’re hitting this error first:
The core issue is that your __call__ method tries to generate an entire sequence of learning rates and iterates over a tf.Tensor (via the list comprehension)—a pattern that’s not allowed in TensorFlow’s graph mode. The LearningRateSchedule interface expects your __call__ method to take a single current step value (automatically passed by Keras during training) and return the corresponding single learning rate for that step—not the full schedule sequence.

Here’s the corrected, graph-compatible implementation of your CustomSchedule class:

import tensorflow as tf
import math

class CustomSchedule(tf.keras.optimizers.schedules.LearningRateSchedule):
    def __init__(self, base_lr=0.1, end_lr=0.001, warmup_steps=390*5, total_steps=390*35):
        super(CustomSchedule, self).__init__()
        self.base_lr = base_lr
        self.end_lr = end_lr
        self.warmup_steps = warmup_steps
        # Store total training steps to calculate cosine decay duration
        self.total_steps = total_steps
        # Precompute the number of steps dedicated to cosine decay
        self.decay_steps = total_steps - warmup_steps

    def __call__(self, step):
        # Convert step to float32 for continuous mathematical operations
        step = tf.cast(step, tf.float32)

        # Warm-up phase: linearly increase LR from 0 to base_lr
        warmup_lr = (self.base_lr / self.warmup_steps) * step

        # Cosine decay phase: smoothly decay from base_lr to end_lr after warmup
        current_decay_step = step - self.warmup_steps
        # Ensure we don't exceed the total decay steps (prevents LR from dropping below end_lr)
        current_decay_step = tf.minimum(current_decay_step, self.decay_steps)
        cosine_lr = self.end_lr + 0.5 * (self.base_lr - self.end_lr) * (
            1 + tf.math.cos(tf.constant(math.pi) * current_decay_step / self.decay_steps)
        )

        # Select the correct LR based on whether we're in warmup or decay phase
        return tf.where(step < self.warmup_steps, warmup_lr, cosine_lr)

Key Fixes & Explanations: #

• Removed sequence generation: Instead of creating full warmup/cosine sequences, we calculate the LR directly for the current step using tensor operations (compatible with graph mode).
• Graph-safe conditionals: Used tf.where instead of Python if/else to ensure the logic works when TensorFlow compiles the training graph.
• Proper step handling: The step parameter passed to __call__ is the live training step (automatically provided by Keras)—we no longer hardcode a default value here.
• Decay bounds: Added tf.minimum to cap the decay step, ensuring the LR doesn’t drop below end_lr once training completes the total planned steps.

How to Use It: #

Initialize your optimizer with an instance of this schedule:

optimizer = tf.keras.optimizers.SGD(learning_rate=CustomSchedule())

You can verify the schedule behaves as expected by plotting it (in eager mode):

import matplotlib.pyplot as plt

schedule = CustomSchedule()
steps = range(390*35)
lrs = [schedule(step).numpy() for step in steps]

plt.plot(steps, lrs)
plt.xlabel("Training Step")
plt.ylabel("Learning Rate")
plt.title("Warm-Up + Cosine Decay Schedule")
plt.show()

This will generate a smooth curve that ramps up linearly during warmup, then decays cosine-wise to your target end LR—all while being fully compatible with TensorFlow’s graph execution mode.

内容的提问来源于stack exchange，提问作者S. P