如何用SessionRunHook实现tf.estimator版DNNLinearCombinedClassifier的精确率/召回率计算

嘿，我来帮你解决这个从ValidationMonitor切换到SessionRunHook计算精确率和召回率的问题！其实只要自定义一个SessionRunHook就能实现相同的功能，下面是具体的思路和代码示例：

tf.estimator.SessionRunHook允许我们在TensorFlow会话的不同生命周期阶段插入自定义逻辑——比如在每次评估步骤后收集真实标签和预测结果，最后统一计算精确率（precision）和召回率（recall）。我们可以通过两种方式实现：一种是借助sklearn的metrics工具，另一种是用TensorFlow原生的指标API。

这种方式代码简洁，适合已经熟悉sklearn的开发者：

import tensorflow as tf
from sklearn.metrics import precision_score, recall_score

class PrecisionRecallHook(tf.estimator.SessionRunHook):
    def __init__(self, labels_tensor, predictions_tensor):
        # 传入真实标签和预测结果的张量
        self.labels_tensor = labels_tensor
        self.predictions_tensor = predictions_tensor
        # 初始化存储结果的列表
        self.true_labels = []
        self.predicted_labels = []

    def before_run(self, run_context):
        # 指定要从会话中获取的张量
        return tf.estimator.SessionRunArgs([self.labels_tensor, self.predictions_tensor])

    def after_run(self, run_context, run_values):
        # 收集每次评估得到的标签和预测值
        labels, preds = run_values.results
        self.true_labels.extend(labels.flatten())
        self.predicted_labels.extend(preds.flatten())

    def end(self, session):
        # 计算并输出指标
        precision = precision_score(self.true_labels, self.predicted_labels)
        recall = recall_score(self.true_labels, self.predicted_labels)
        print(f"评估完成 - 精确率: {precision:.4f}, 召回率: {recall:.4f}")
        # 也可以把结果写入文件或TensorBoard，按需扩展

如果想完全贴合TensorFlow生态，不依赖外部库，可以用TF原生的tf.metrics模块：

import tensorflow as tf

class PrecisionRecallTFHook(tf.estimator.SessionRunHook):
    def __init__(self, labels_tensor, predictions_tensor):
        self.labels_tensor = labels_tensor
        self.predictions_tensor = predictions_tensor
        # 初始化精确率和召回率的指标变量
        self.precision, self.update_precision = tf.metrics.precision(labels_tensor, predictions_tensor)
        self.recall, self.update_recall = tf.metrics.recall(labels_tensor, predictions_tensor)

    def begin(self):
        # 定义指标变量的初始化操作
        self.init_local_vars = tf.group(tf.local_variables_initializer())

    def after_create_session(self, session, coord):
        # 会话创建后，初始化本地变量
        session.run(self.init_local_vars)

    def before_run(self, run_context):
        # 指定每次评估要运行的指标更新操作
        return tf.estimator.SessionRunArgs([self.update_precision, self.update_recall])

    def end(self, session):
        # 获取最终的指标结果并输出
        precision_val, recall_val = session.run([self.precision, self.recall])
        print(f"评估完成 - 精确率: {precision_val:.4f}, 召回率: {recall_val:.4f}")

接下来需要在你的model_fn中，把这个Hook关联到评估模式下的EstimatorSpec：

def custom_model_fn(features, labels, mode, params):
    # 初始化DNNLinearCombinedClassifier的基础模型
    base_classifier = tf.estimator.DNNLinearCombinedClassifier(
        linear_feature_columns=params['linear_columns'],
        dnn_feature_columns=params['dnn_columns'],
        dnn_hidden_units=[128, 64],
        n_classes=2,  # 根据你的任务调整类别数
        model_dir='./your_model_dir'
    )

    # 调用基础分类器的model_fn获取EstimatorSpec
    estimator_spec = base_classifier.model_fn(features, labels, mode, params)

    if mode == tf.estimator.ModeKeys.EVAL:
        # 获取预测的类别（二分类取logits的argmax，多分类按需调整）
        predicted_classes = tf.argmax(estimator_spec.predictions['logits'], axis=1)
        # 实例化自定义Hook
        pr_hook = PrecisionRecallHook(labels, predicted_classes)  # 或者用PrecisionRecallTFHook
        # 将Hook添加到评估钩子列表中
        estimator_spec = estimator_spec._replace(
            evaluation_hooks=[pr_hook]
        )

    return estimator_spec

# 创建Estimator并执行评估
estimator = tf.estimator.Estimator(model_fn=custom_model_fn, params=your_params)
eval_results = estimator.evaluate(input_fn=your_eval_input_fn)

• 如果是多分类任务，记得调整指标计算的参数：比如sklearn的precision_score要指定average='macro'或'weighted'；TF原生的tf.metrics.precision可以设置class_id来计算特定类别的指标。
• 如果你想把指标写入TensorBoard，可以在Hook的end方法中添加tf.summary操作，或者直接把指标加入estimator_spec的eval_metric_ops中（不过这样Estimator会自动输出这些指标，不需要Hook了，但如果需要更自定义的逻辑，Hook还是更灵活）。

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