基于Stratified KFold的sklearn模型多测试集拆分预测生成与整合报错排查及解决方案问询

我来帮你搞定这个问题！你的代码里有几个关键细节没处理好，导致了ValueError，我会一步步修正，最终得到包含原数据行索引和对应预测概率的结果。

问题分析 #

你遇到的核心问题有这几个：

• train_model函数里错误地使用了全局的X和y（整个数据集的特征和目标）去索引训练/测试子集，这会触发索引错误
• 传入模型的y_train是DataFrame格式，但sklearn模型要求目标变量是一维的Series/数组
• 没有保留测试集的原始索引，导致最后无法将预测结果对应到原数据的具体行
• 直接用pd.concat合并数组类型的结果，无法保留索引信息

修正后的完整代码 #

import pandas as pd
from sklearn import datasets
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import StratifiedKFold

# 加载数据，组装DataFrame
iris = datasets.load_iris()
X = pd.DataFrame(iris.data[51:150, :], columns=["sepal_length", "sepal_width", "petal_length", "petal_width"])
y = pd.DataFrame(iris.target[51:150,], columns=["target"])
df = pd.concat([X, y], axis=1)

# 实例化逻辑回归模型（增加max_iter避免收敛警告）
log = LogisticRegression(max_iter=200)

# 修正后的建模函数
def train_model(train_df, test_df):
    # 从传入的训练/测试子集里提取特征和目标
    X_train = train_df.drop("target", axis=1)
    y_train = train_df["target"]  # 转成Series，符合sklearn要求
    X_test = test_df.drop("target", axis=1)
    
    # 训练模型并生成类别1的预测概率，同时保留测试集的索引
    prob = log.fit(X_train, y_train).predict_proba(X_test)[:, 1]
    # 将结果转成带索引的Series，方便后续合并
    return pd.Series(prob, index=test_df.index, name="pred_prob")

# 生成分层K折拆分（这里用2折做示例，加shuffle保证可复现）
skf = StratifiedKFold(n_splits=2, shuffle=True, random_state=42)

# 生成并整合所有预测结果
outputs = []
for train_index, test_index in skf.split(df, df["target"]):
    train_df = df.loc[train_index, :]
    test_df = df.loc[test_index, :]
    fold_pred = train_model(train_df, test_df)
    outputs.append(fold_pred)

# 合并所有折叠的预测结果，按原索引排序
all_preds = pd.concat(outputs).sort_index()
print(all_preds.head())

关键修改点说明 #

1. 修改train_model函数：

   • 不再依赖全局的X和y，而是直接从传入的train_df和test_df中提取特征和目标
   • y_train改用train_df["target"]得到一维Series，满足sklearn模型的输入要求
   • 将预测概率转成带测试集索引的Series，确保每个概率能对应到原数据的行

2. 优化交叉验证循环：

   • 给StratifiedKFold加上shuffle=True和random_state，保证拆分的可复现性
   • 收集每个折叠的带索引Series结果，而不是单纯的数组

3. 合并结果：

   • 用pd.concat合并Series后，调用sort_index()让结果和原数据的行顺序一致

运行这段代码后，all_preds就是一个包含原数据行索引和对应类别1预测概率的Series，你可以直接把它和原df合并使用。

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