使用GroupKFold调LGBMRanker参数报Sum of query counts错误如何解决

错误原因 #

你向RandomizedSearchCV的fit方法传入的group参数是全局训练集的query计数数组qids_train，它的总和等于完整X_train的样本数。但交叉验证运行时，每折只会取X_train的子集作为折内训练数据，LightGBM拿到的还是全局的qids_train，总和和折内样本数不匹配，就触发了该报错。

另外你原代码里用普通train_test_split拆分数据集对排序任务是错误的：普通拆分可能把同一个query的样本分别拆到训练集和验证集，排序任务要求同一个query的所有样本必须属于同一个集合，不能跨集合拆分。

解决方案 #

直接修改参数传递逻辑即可，LGBMRanker的fit方法的group参数支持直接传入每个样本对应的query_id值，不需要提前聚合为每个query的样本数，它会自动根据当前传入的样本批次实时计算query计数，完美适配交叉验证的折内子集场景。

修改后的完整代码：

# 替换原train_test_split，按query分组拆分训练验证集
from sklearn.model_selection import GroupShuffleSplit
gss = GroupShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
train_idx, val_idx = next(gss.split(X, y, groups=X["query_id"]))
X_train, X_val = X.iloc[train_idx], X.iloc[val_idx]
y_train, y_val = y.iloc[train_idx], y.iloc[val_idx]

gbm = lightgbm.LGBMRanker(
    objective="lambdarank",
    metric="ndcg",
    is_unbalance=True,
)

gkf = GroupKFold(n_splits=5)
# 直接用query_id作为分组依据，无需自行生成flatted_group_train
cv = gkf.split(X_train, y_train, groups=X_train["query_id"])
grid = RandomizedSearchCV(gbm, params_grid, n_iter=10, cv=cv, verbose=2, refit=False)
# 传入每个样本对应的query_id作为group参数，不要传聚合后的qids_train
grid.fit(
    X=X_train,
    y=y_train,
    group=X_train["query_id"],
)

可选校验 #

如果修改后仍报错，可以加以下代码提前校验数据合法性：

# 校验全局训练集的query计数总和是否和样本数一致
assert len(X_train) == X_train.groupby("query_id")["query_id"].count().sum(), "训练集query计数总和与样本数不匹配"
# 校验无query仅含单个样本的情况（LightGBM排序默认要求每个query至少2个样本）
assert (X_train.groupby("query_id")["query_id"].count() >= 2).all(), "存在仅含1个样本的query，会导致训练报错"

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