如何在TensorFlow Estimator中获取各类别的准确率

嘿，我来帮你搞定这个问题——默认的classifier.evaluate()只会返回整体的准确率，要拿到每个类别的准确率，其实有两种实用的路径，我给你一步步拆解：

方法一：通过预测结果+混淆矩阵计算 #

这种方法很直观，先拿到所有测试样本的真实标签和模型的预测标签，再用混淆矩阵计算每个类别的准确率（每个类别的准确率=该类别正确预测数/该类别总样本数）。

步骤1：收集真实标签与预测标签 #

先准备好你的测试集输入函数，然后遍历模型的预测结果，同时收集真实标签：

# 假设你已经定义好了测试集输入函数test_input_fn
predictions = classifier.predict(input_fn=test_input_fn)

# 初始化列表存储真实标签和预测标签
true_labels = []
predicted_labels = []

# 这里的test_data需要是你的测试数据集（包含真实标签字段，比如label）
for pred, sample in zip(predictions, test_data):
    true_labels.append(sample['label'])
    # 从预测结果中取出类别ID
    predicted_labels.append(pred['class_ids'][0])

步骤2：计算混淆矩阵与类别准确率 #

借助sklearn的混淆矩阵工具，快速计算每个类别的准确率：

from sklearn.metrics import confusion_matrix
import numpy as np

# 生成混淆矩阵：行是真实类别，列是预测类别
confusion_mat = confusion_matrix(true_labels, predicted_labels)

# 计算每个类别的准确率：对角线是正确数，行求和是该类别总样本数
class_accuracies = np.diag(confusion_mat) / confusion_mat.sum(axis=1)

# 输出结果就是你想要的数组形式
print(class_accuracies)
# 示例输出：[0.924012, 0.89478, 0.945123, 0.685042]

方法二：自定义评估指标集成到Estimator #

如果你想让classifier.evaluate()直接返回每个类别的准确率，可以自定义评估指标函数，把它集成到Estimator的配置里：

步骤1：定义自定义评估函数 #

这个函数会针对每个类别计算准确率，处理好除以0的边界情况：

import tensorflow as tf

def class_wise_accuracy_metrics(labels, predictions):
    # 替换成你的类别总数，比如这里是4类
    num_classes = 4
    metric_dict = {}
    
    for class_idx in range(num_classes):
        # 标记该类别下的真实样本和正确预测的样本
        is_true_class = tf.cast(tf.equal(labels, class_idx), tf.float32)
        is_correct_pred = tf.cast(
            tf.logical_and(tf.equal(labels, class_idx), 
                          tf.equal(predictions['class_ids'], class_idx)),
            tf.float32
        )
        
        # 计算该类别的准确率，避免除以0
        class_accuracy = tf.divide(
            tf.reduce_sum(is_correct_pred),
            tf.reduce_sum(is_true_class) + 1e-8,  # 加小值防止除以0
            name=f'class_{class_idx}_accuracy'
        )
        
        metric_dict[f'class_{class_idx}_accuracy'] = class_accuracy
    
    return metric_dict

步骤2：创建Estimator时添加自定义指标 #

在初始化DNNClassifier（或其他Premade Estimator）时，把自定义指标传入eval_metric_ops参数：

classifier = tf.estimator.DNNClassifier(
    feature_columns=your_feature_columns,  # 替换成你的特征列
    hidden_units=[1024, 512, 256],  # 替换成你的网络结构
    n_classes=4,  # 替换成你的类别总数
    model_dir='./your_model_dir',  # 替换成你的模型保存路径
    eval_metric_ops=class_wise_accuracy_metrics  # 加入自定义评估指标
)

步骤3：执行评估并获取结果 #

现在调用evaluate()后，返回的结果字典里就会包含每个类别的准确率：

eval_result = classifier.evaluate(input_fn=test_input_fn)

# 提取每个类别的准确率组成数组
class_accuracies = [eval_result[f'class_{i}_accuracy'] for i in range(4)]
print(class_accuracies)

两种方法各有优势：方法一适合快速验证，不需要修改Estimator的初始化；方法二更贴合TensorFlow的Estimator流程，适合正式的训练评估流水线。

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