如何让主代码的Stratified-K-fold与learning_curve例程使用相同折分？

嘿，这个问题我之前踩过坑！其实解决起来很简单——核心就是让learning_curve复用你主逻辑里已经定义好的分层折分，而不是让它自己重新生成一套。下面是具体的实现步骤：

1. 先定义统一的交叉验证拆分器 #

首先创建一个参数固定的StratifiedKFold实例，指定好折数、是否洗牌、随机种子（保证拆分可复现），这样不管是主训练逻辑还是学习曲线绘制，都会遵循同一套拆分规则。

from sklearn.model_selection import StratifiedKFold
import numpy as np

# 根据你的需求调整参数，比如n_splits是折数，random_state保证拆分固定
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

2. 主训练逻辑中使用该拆分器 #

在你的主体代码里，直接用这个skf实例来拆分数据，循环每个fold完成训练和评估：

# 假设X是特征矩阵，y是标签数据
for train_idx, val_idx in skf.split(X, y):
    # 通过索引拆分训练集与验证集
    X_train, X_val = X[train_idx], X[val_idx]
    y_train, y_val = y[train_idx], y[val_idx]
    
    # 你的模型训练、评估逻辑
    your_model.fit(X_train, y_train)
    train_acc = your_model.score(X_train, y_train)
    val_acc = your_model.score(X_val, y_val)
    
    # 记录得分、保存模型等操作...

3. 调用learning_curve时传入同一个拆分器 #

关键一步：调用learning_curve的时候，把刚才定义好的skf实例直接传给cv参数，而不是使用默认的自动生成逻辑。这样learning_curve就会复用你已经定义好的所有折分：

from sklearn.model_selection import learning_curve
import matplotlib.pyplot as plt

# 生成学习曲线所需的得分数据
train_sizes, train_scores, val_scores = learning_curve(
    estimator=your_model,  # 你的模型实例
    X=X,
    y=y,
    cv=skf,  # 核心：传入同一组StratifiedKFold拆分器
    train_sizes=np.linspace(0.1, 1.0, 10),  # 训练集大小的比例区间
    scoring='accuracy'  # 根据任务选择评估指标，比如'f1'、'roc_auc'等
)

# 绘制学习曲线（示例代码）
plt.figure(figsize=(10,6))
plt.plot(train_sizes, np.mean(train_scores, axis=1), label='Training Score')
plt.plot(train_sizes, np.mean(val_scores, axis=1), label='Validation Score')
plt.xlabel('Training Set Size')
plt.ylabel('Score')
plt.title('Learning Curve')
plt.legend()
plt.show()

额外技巧：提前保存折分索引 #

如果你想彻底固定折分（避免多次调用split可能出现的意外），还可以提前把所有折分的索引对保存成列表，然后在主逻辑和learning_curve里都使用这个列表：

# 预先生成并保存所有折分的索引对
fold_indices = list(skf.split(X, y))

# 主逻辑中使用该列表
for train_idx, val_idx in fold_indices:
    # ...训练逻辑

# learning_curve中同样传入该列表
train_sizes, train_scores, val_scores = learning_curve(
    estimator=your_model,
    X=X,
    y=y,
    cv=fold_indices,  # 直接传预先生成的索引对列表
    # 其他参数...
)

这样操作后，你的主训练逻辑和学习曲线绘制就会使用完全一致的训练/验证拆分，避免了因折分不同导致的得分差异，保证了结果的可比性。

内容的提问来源于stack exchange，提问作者Fränzu