TensorFlow决策森林自定义PR-AUC指标随树数量变化的获取方法

你可以利用TFDF树集成模型原生支持的截断推理能力实现需求，不需要修改训练逻辑，仅需在训练完成后按树数量逐层截断预测并计算指标即可，具体实现如下：

1. 按你现有逻辑完成模型训练
2. 遍历1到总树数的取值，调用predict时通过inference_args参数指定仅使用前N棵树推理，计算对应PR-AUC值

完整代码示例（接你现有训练完成的代码）：

# 获取模型总树数
inspector = tfdf_clf.make_inspector()
total_trees = inspector.num_trees
# 存储每棵树数量对应的PR-AUC结果
pr_auc_curve = []

for n_tree in range(1, total_trees + 1):
    # 仅用前n_tree棵树做预测
    y_pred = tfdf_clf.predict(test_ds, verbose=0, inference_args={"num_trees": n_tree})
    # 重置并计算PR-AUC
    pr_auc.reset_state()
    for batch_x, batch_y in test_ds:
        pr_auc.update_state(batch_y, y_pred)
    pr_auc_curve.append(pr_auc.result().numpy())

# pr_auc_curve第i个元素即为使用i+1棵树时的测试集PR-AUC
print(pr_auc_curve)

补充说明 #

• 该方法仅需训练一次模型，通过推理阶段截断的方式获取不同树数量的结果，计算效率很高
• inference_args是TFDF Keras模型原生支持的参数，符合仅使用TFDF库的要求
• 如果需要同时获取训练集的指标变化，仅需将test_ds替换为train_ds即可

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