如何让tf.estimator.LinearClassifier.train()显示类似model.fit的训练进度？

需求说明 #

使用tf.estimator.LinearClassifier的train()方法时，希望实现与model.fit()一致的训练进度输出，每个Epoch显示步数、耗时、损失、准确率等信息，期望输出样式如下：

Epoch 1/10
1875/1875 [==============================] - 5s 3ms/step - loss: 0.4964 - accuracy: 0.8270
Epoch 2/10
1875/1875 [==============================] - 4s 2ms/step - loss: 0.3751 - accuracy: 0.8652
Epoch 3/10
1875/1875 [==============================] - 5s 3ms/step - loss: 0.3382 - accuracy: 0.8762

原示例代码：

# 输入函数
def make_input_fn(data_df, label_df, num_epochs=1000, shuffle=True, batch_size=32):
  def input_function():  # 内部函数，最终返回该函数
    ds = tf.data.Dataset.from_tensor_slices((dict(data_df), label_df))  # 创建包含数据和标签的tf.data.Dataset对象
    if shuffle:
      ds = ds.shuffle(1000)  # 随机打乱数据顺序
    ds = ds.batch(batch_size).repeat(num_epochs)  # 按batch_size拆分数据集，并重复指定epoch次数
    return ds  # 返回一个批次的数据集
  return input_function  # 返回供后续使用的函数对象

train_input_fn = make_input_fn(dftrain, y_train)  # 调用返回的input_function获取可传入模型的数据集对象
eval_input_fn = make_input_fn(dfeval, y_eval, num_epochs=1, shuffle=False)
pre_input_fn = make_input_fn(dfpre, y_pre, num_epochs=1, shuffle=False)

linear_est = tf.estimator.LinearClassifier(feature_columns=feature_columns)
 
linear_est.train(train_input_fn)  # 训练
result = linear_est.evaluate(eval_input_fn)

实现方案 #

tf.estimator默认不提供model.fit()风格的进度输出，需要自定义SessionRunHook来捕获训练过程中的指标、记录耗时，并在每个Epoch结束时格式化输出。具体步骤如下：

1. 计算单Epoch步数：根据训练集总样本数和batch_size，计算每个Epoch包含的训练步数。
2. 自定义训练Hook：通过tf.estimator.SessionRunHook在训练开始、每个step结束、Epoch结束时执行对应的逻辑，记录耗时、损失、准确率等指标。
3. 传入Hook到train方法：调用linear_est.train()时，将自定义Hook传入hooks参数。

修改后的完整代码 #

import tensorflow as tf
import time

# 自定义训练进度Hook
class TrainingProgressHook(tf.estimator.SessionRunHook):
    def __init__(self, num_steps_per_epoch, num_epochs):
        self.num_steps_per_epoch = num_steps_per_epoch
        self.num_epochs = num_epochs
        self.current_epoch = 0
        self.current_step = 0
        self.epoch_start_time = None
        self.total_loss = 0.0
        self.total_accuracy = 0.0

    def begin(self):
        # 训练开始前初始化
        pass

    def before_run(self, run_context):
        if self.current_step % self.num_steps_per_epoch == 0:
            self.epoch_start_time = time.time()
            self.current_epoch += 1
            print(f"\nEpoch {self.current_epoch}/{self.num_epochs}")
            self.total_loss = 0.0
            self.total_accuracy = 0.0
        # 指定要获取的张量（损失和准确率，需与模型的输出对应）
        return tf.estimator.SessionRunArgs({
            "loss": tf.compat.v1.losses.get_total_loss(),
            "accuracy": tf.compat.v1.metrics.accuracy()[1]
        })

    def after_run(self, run_context, run_values):
        self.current_step += 1
        step_loss = run_values.results["loss"]
        step_accuracy = run_values.results["accuracy"]
        self.total_loss += step_loss
        self.total_accuracy += step_accuracy

        # 每个Epoch结束时输出汇总信息
        if self.current_step % self.num_steps_per_epoch == 0:
            epoch_duration = time.time() - self.epoch_start_time
            avg_loss = self.total_loss / self.num_steps_per_epoch
            avg_accuracy = self.total_accuracy / self.num_steps_per_epoch
            step_time = (epoch_duration / self.num_steps_per_epoch) * 1000  # 转换为毫秒

            # 格式化输出进度条
            progress_bar = "[" + "="*30 + "]"
            print(f"{self.num_steps_per_epoch}/{self.num_steps_per_epoch} {progress_bar} - {int(epoch_duration)}s {int(step_time)}ms/step - loss: {avg_loss:.4f} - accuracy: {avg_accuracy:.4f}")

# 输入函数
def make_input_fn(data_df, label_df, num_epochs=1000, shuffle=True, batch_size=32):
    def input_function():
        ds = tf.data.Dataset.from_tensor_slices((dict(data_df), label_df))
        if shuffle:
            ds = ds.shuffle(1000)
        ds = ds.batch(batch_size).repeat(num_epochs)
        return ds
    return input_function

# 假设dftrain、y_train、feature_columns已定义
train_samples = len(dftrain)
batch_size = 32
num_epochs = 10
num_steps_per_epoch = train_samples // batch_size

train_input_fn = make_input_fn(dftrain, y_train, num_epochs=num_epochs, batch_size=batch_size)
eval_input_fn = make_input_fn(dfeval, y_eval, num_epochs=1, shuffle=False)
pre_input_fn = make_input_fn(dfpre, y_pre, num_epochs=1, shuffle=False)

linear_est = tf.estimator.LinearClassifier(feature_columns=feature_columns)

# 创建自定义Hook并传入train方法
progress_hook = TrainingProgressHook(num_steps_per_epoch, num_epochs)
linear_est.train(train_input_fn, hooks=[progress_hook])

result = linear_est.evaluate(eval_input_fn)

注意事项 #

• 确保num_steps_per_epoch的计算准确：如果训练集样本数不能被batch_size整除，可根据需求选择向上取整或忽略剩余样本。
• 张量名称需匹配：before_run中获取的损失和准确率张量，要与LinearClassifier内部计算的张量一致，若有自定义指标需调整对应张量。
• 该Hook仅适配TensorFlow 2.x版本的tf.estimator兼容模式，确保环境为TensorFlow 2.9.2。

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