使用Estimator训练自定义Keras模型时如何打印指标值？

我之前也碰到过一模一样的情况——用model_to_estimator转完模型后，训练日志里只有loss，自己定义的metrics死活不显示，哪怕加了add_metrics也没用。其实问题出在Estimator默认的日志逻辑和Keras模型转换后的细节上，给你几个实用的解决办法：

方法1：用LoggingTensorHook强制打印训练指标 #

当你把Keras模型转成Estimator后，Keras定义的metrics（比如你用的accuracy）其实已经存在于TensorFlow的计算图里了，只是Estimator默认不会在训练阶段主动打印这些指标。你可以添加一个日志钩子，指定要打印的张量：

# 先完成模型转换
estimator_model = tf.keras.estimator.model_to_estimator(
    keras_model=model, model_dir=FLAGS.model_path)

# 定义日志钩子，指定要打印的张量名称
# 注意：张量名称要和计算图里的实际名称对应，比如accuracy可能叫'acc'或者'accuracy'
logging_hook = tf.train.LoggingTensorHook(
    tensors={'训练loss': 'loss', '训练准确率': 'acc'},
    every_n_iter=100)  # 每100步打印一次

# 训练时带上这个钩子
estimator_model.train(
    input_fn=input_fn,
    hooks=[logging_hook])

如果不确定张量的准确名称，可以用tf.get_default_graph().get_all_tensor_names()打印所有张量名称，找到对应的metrics张量。

方法2：正确使用tf.estimator.add_metrics #

你之前用add_metrics无效，大概率是自定义的metric函数参数不对。这个函数要求metric必须接收features和labels作为输入，并且返回(metric_value, update_op)的元组。比如你要加AUC的话，得这么写：

def my_auc(features, labels):
    # 获取模型的预测结果
    predictions = estimator_model.model_fn(features, None, tf.estimator.ModeKeys.TRAIN)[0]['predictions']
    # 计算AUC和对应的更新操作
    auc, auc_update_op = tf.metrics.auc(labels, predictions)
    return auc, auc_update_op

# 给Estimator添加这个metric
estimator_model = tf.estimator.add_metrics(estimator_model, my_auc)

不过要注意，add_metrics添加的指标默认只会在评估阶段显示，如果想在训练时打印，还是得配合上面的LoggingTensorHook。

方法3：自定义Estimator的model_fn（最灵活） #

如果前两个方法都不好用，不如直接手动写Estimator的model_fn，把Keras模型包进去，这样能完全掌控metrics的计算和日志输出：

def model_fn(features, labels, mode):
    # 实例化你的自定义Keras模型
    model = EstimatorModel()
    # 获取模型输出
    logits = model(features)
    
    # 计算损失
    loss = tf.keras.losses.categorical_crossentropy(labels, logits)
    loss = tf.reduce_mean(loss)
    
    # 计算准确率和AUC指标
    accuracy = tf.metrics.accuracy(
        tf.argmax(labels, axis=1), tf.argmax(logits, axis=1))
    auc = tf.metrics.auc(
        tf.argmax(labels, axis=1), tf.argmax(logits, axis=1))
    
    # 根据模式返回EstimatorSpec
    if mode == tf.estimator.ModeKeys.TRAIN:
        optimizer = tf.keras.optimizers.Adam()
        train_op = optimizer.get_updates(loss, model.trainable_variables)
        # 训练时添加日志钩子打印指标
        train_hook = tf.train.LoggingTensorHook(
            {'loss': loss, 'accuracy': accuracy[0]}, every_n_iter=100)
        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            train_op=train_op,
            training_hooks=[train_hook])
    elif mode == tf.estimator.ModeKeys.EVAL:
        # 评估阶段返回指标
        return tf.estimator.EstimatorSpec(
            mode=mode,
            loss=loss,
            eval_metric_ops={'accuracy': accuracy, 'auc': auc})

# 创建自定义Estimator
estimator_model = tf.estimator.Estimator(model_fn=model_fn, model_dir=FLAGS.model_path)
estimator_model.train(input_fn=input_fn)

最后还要提醒你几个细节：

• 确保你的input_fn返回的labels是one-hot编码的，毕竟你用的是categorical_crossentropy损失；
• 检查TensorFlow版本，你的代码看起来是TF1.x的风格，不同版本的API可能有细微差异。

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