如何在TensorFlow决策森林GBDT排序任务训练中分步记录自定义评估指标

方法一：分批次渐进训练+每轮指标评估 #

TF-DF的GradientBoostedTreesModel支持通过num_trees参数控制累计训练的树数量，我们可以分阶段训练，每完成一批树的训练后就暂停，手动计算并记录自定义验证指标，再继续训练剩余的树。

步骤如下：

1. 初始化模型时设置初始树数量为0，完成基础配置。
2. 循环执行「训练N棵树 → 评估自定义指标」的流程，每次训练指定累计要达到的树总数，确保训练进度延续。

示例代码：

import tensorflow_decision_forests as tfdf
import pandas as pd
import numpy as np
from sklearn.metrics import ndcg_score

# 加载训练/验证数据，确保包含排序分组列（如query_id）
train_df = pd.read_csv("train_data.csv")
val_df = pd.read_csv("val_data.csv")

train_ds = tfdf.keras.pd_dataframe_to_tf_dataset(
    train_df, label="relevance", task=tfdf.keras.Task.RANKING
)
val_ds = tfdf.keras.pd_dataframe_to_tf_dataset(
    val_df, label="relevance", task=tfdf.keras.Task.RANKING
)

# 初始化模型，初始树数量设为0
model = tfdf.keras.GradientBoostedTreesModel(
    task=tfdf.keras.Task.RANKING,
    num_trees=0,
    ranking_group="query_id"  # 排序任务必须指定分组列
)
model.compile()

# 自定义指标：加权NDCG（以真实相关性为权重）
def custom_weighted_ndcg(y_true, y_pred, query_ids):
    unique_queries = np.unique(query_ids)
    scores = []
    for qid in unique_queries:
        mask = query_ids == qid
        true_scores = y_true[mask]
        pred_scores = y_pred[mask]
        # 用真实相关性作为权重计算NDCG
        weighted_true = true_scores * true_scores
        scores.append(ndcg_score([weighted_true], [pred_scores]))
    return np.mean(scores)

# 分阶段训练参数
total_trees = 100
step_trees = 10  # 每轮训练10棵树
current_trees = 0

# 循环训练并评估
while current_trees < total_trees:
    current_trees += step_trees
    # 训练到累计current_trees棵树（延续之前的训练进度）
    model.fit(train_ds, num_trees=current_trees)
    
    # 提取验证集的标签、预测值和分组ID
    val_labels, val_preds, val_query_ids = [], [], []
    for batch in val_ds:
        features, labels = batch
        preds = model.predict(features, verbose=0)
        val_labels.extend(labels.numpy().flatten())
        val_preds.extend(preds.flatten())
        val_query_ids.extend(features["query_id"].numpy().flatten())
    
    # 计算并打印自定义指标
    metric_value = custom_weighted_ndcg(
        np.array(val_labels), np.array(val_preds), np.array(val_query_ids)
    )
    print(f"已训练{current_trees}棵树，自定义加权NDCG: {metric_value:.4f}")

方法二：训练完成后回溯计算每步指标 #

如果不需要实时监控，可先完成完整训练，再利用模型的Inspector工具获取训练日志，回溯每棵树训练后的模型状态，计算验证集的自定义指标：

1. 完整训练模型后，用model.make_inspector().training_logs()获取每棵树的训练日志。
2. 通过model._model.n_trees()手动控制模型使用的树数量，模拟每一步的训练状态，计算对应阶段的验证指标。

这种方法适合事后分析训练过程中指标的变化趋势。

关键注意事项 #

• fit()方法中的num_trees是累计值而非增量值，必须指定累计要达到的树总数才能延续训练进度。
• 排序任务的自定义指标必须按分组（如query_id）计算，不能直接对全局样本求值。
• 验证集的数据集格式要与训练集一致，确保分组列和标签列正确传递。

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