训练tf.estimator时记录准确率指标及预构建Estimator打印指标最简方法

嘿，我明白你纠结的点——用预构建Estimator的时候，不想折腾自定义那套冗余代码，又想在训练时记录并打印准确率和损失对吧？这就给你说最简的解决办法：

对于预构建Estimator（比如DNNClassifier、LinearClassifier这类），不用重写model_fn，只需要用tf.estimator.add_metrics给它追加准确率指标就行，步骤如下：

• 先定义一个指标函数，接收模型的预测结果和标签，返回你要记录的准确率指标：

def accuracy_metric_fn(predictions, labels):
    # 预构建Estimator的predictions是个字典，class_ids对应分类的预测类别
    return {
        'train_accuracy': tf.metrics.accuracy(
            labels=labels,
            predictions=predictions['class_ids']
        )
    }

• 创建你的预构建Estimator后，调用tf.estimator.add_metrics把指标绑定上去：

# 举个例子，创建一个DNN分类器
classifier = tf.estimator.DNNClassifier(
    feature_columns=feature_cols,
    hidden_units=[128, 64],
    n_classes=10
)

# 追加准确率指标
tf.estimator.add_metrics(classifier, accuracy_metric_fn)

这样一来，训练过程中这个准确率指标就会被自动记录到TensorBoard（如果你指定了模型保存路径的话），同时也能通过Hook来打印。

最简方式就是用tf.train.LoggingTensorHook，指定要打印的张量名称，然后把这个Hook传给训练函数。

• 首先确定要打印的张量名称：

  • 损失张量的默认名称是loss（预构建Estimator自带）
  • 我们刚才追加的准确率张量名称是train_accuracy（就是指标函数里定义的key）

• 创建LoggingHook：

logging_hook = tf.train.LoggingTensorHook(
    tensors={'loss': 'loss', 'accuracy': 'train_accuracy'},
    every_n_iter=100  # 每100步打印一次
)

• 最后用tf.estimator.train或者tf.estimator.train_and_evaluate启动训练，把Hook传进去：

# 用train的方式
classifier.train(
    input_fn=train_input_fn,
    steps=1000,
    hooks=[logging_hook]
)

# 或者用train_and_evaluate（更推荐，支持自动评估）
train_spec = tf.estimator.TrainSpec(
    input_fn=train_input_fn,
    max_steps=1000,
    hooks=[logging_hook]
)
eval_spec = tf.estimator.EvalSpec(input_fn=eval_input_fn)
tf.estimator.train_and_evaluate(classifier, train_spec, eval_spec)

这样每训练100步，控制台就会输出类似这样的内容：

INFO:tensorflow:loss = 0.345, accuracy = 0.92, step = 100
INFO:tensorflow:loss = 0.210, accuracy = 0.95, step = 200

补充说明 #

• 如果你找不到张量的准确名称，可以在训练前先运行一次模型的predict或者train一步，然后用classifier.get_variable_names()或者查看TensorBoard的图结构来确认。
• 预构建Estimator本身已经包含损失张量，所以不需要额外定义，直接用loss就行。
• 这种方式完全不需要自定义Estimator，完美适配现成的预构建模型，避免冗余代码。

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