如何在TensorFlow Estimators中跨epoch调整优化器学习率

嘿，我明白你想用Estimators API（尤其是DNNClassifier）来在不同epoch之间调整学习率，而不是用传统的sess.run+placeholder那套。其实Estimators里有几种优雅的方式来实现这个，我给你拆解一下：

TensorFlow提供了现成的学习率衰减函数（比如阶梯衰减、指数衰减），你可以直接把动态计算的学习率传给优化器，Estimator会自动在每个step/epoch帮你更新它。

比如阶梯衰减（每N个epoch就降低学习率）的代码示例：
首先定义学习率的计算逻辑：

def get_learning_rate():
    # 获取当前全局步数
    global_step = tf.train.get_global_step()
    # 初始学习率
    initial_lr = 0.1
    # 假设每个epoch有1000个step，每10个epoch衰减一次
    decay_steps = 10 * 1000
    # 定义不同阶段的学习率
    lr = tf.train.piecewise_constant_decay(
        global_step,
        boundaries=[decay_steps, 2*decay_steps],
        values=[initial_lr, initial_lr*0.5, initial_lr*0.25]
    )
    return lr

然后在创建DNNClassifier时，通过optimizer参数传入带动态学习率的优化器：

classifier = tf.estimator.DNNClassifier(
    feature_columns=feature_columns,
    hidden_units=[1024, 512, 256],
    # 用lambda包裹，确保模型构建时能拿到最新的学习率
    optimizer=lambda: tf.train.GradientDescentOptimizer(learning_rate=get_learning_rate()),
    n_classes=10
)

如果你需要更灵活的控制（比如根据验证集精度调整、手动指定特定epoch调整），可以用SessionRunHook来实现：

1. 先定义一个Hook类，用于在epoch结束时更新学习率：

class LearningRateSchedulerHook(tf.train.SessionRunHook):
    def __init__(self, lr_var, schedule):
        self.lr_var = lr_var  # 保存学习率的变量
        # 调度规则：key是epoch数，value是对应学习率
        self.schedule = schedule
        self.current_epoch = -1

    def begin(self):
        # 通过全局步数计算当前epoch（需提前知道每个epoch的step数）
        self.epoch_tensor = tf.compat.v1.train.get_or_create_global_step() // steps_per_epoch

    def before_run(self, run_context):
        # 每次run前获取当前epoch
        return tf.train.SessionRunArgs(self.epoch_tensor)

    def after_run(self, run_context, run_values):
        epoch = run_values.results
        # 如果进入新epoch且该epoch有调度规则，更新学习率
        if epoch != self.current_epoch and epoch in self.schedule:
            new_lr = self.schedule[epoch]
            run_context.session.run(self.lr_var.assign(new_lr))
            print(f"Epoch {epoch}: 更新学习率为 {new_lr}")
            self.current_epoch = epoch

2. 然后创建模型并传入Hook：

# 定义可更新的学习率变量（trainable=False避免被优化）
initial_lr = 0.1
lr_var = tf.Variable(initial_lr, trainable=False)

# 自定义调度规则：第0、5、10 epoch对应不同学习率
lr_schedule = {0:0.1, 5:0.05, 10:0.01}
# 计算每个epoch的step数：训练集总数 // batch_size
steps_per_epoch = len(train_data) // batch_size

# 创建学习率调度Hook
lr_hook = LearningRateSchedulerHook(lr_var, lr_schedule)

# 创建DNNClassifier
classifier = tf.estimator.DNNClassifier(
    feature_columns=feature_columns,
    hidden_units=[1024, 512, 256],
    optimizer=lambda: tf.train.GradientDescentOptimizer(learning_rate=lr_var),
    n_classes=10
)

# 训练时传入Hook
classifier.train(
    input_fn=train_input_fn,
    steps=total_steps,
    hooks=[lr_hook]
)

• 如果你用的是TensorFlow 2.x，建议优先使用Keras Estimator，API会更统一
• 方法1需要预先计算每个epoch的step数，确保衰减节点设置准确
• 方法2的epoch计算依赖全局步数，如果你用tf.data.Dataset的repeat()，这个逻辑是成立的

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