python - 使用 StratifiedKFold 创建训练/测试/验证拆分

Question

我正在尝试使用 StratifiedKFold 创建用于非 sklearn 机器学习工作流程的训练/测试/验证拆分。因此，DataFrame 需要拆分然后保持原样。

我正在尝试像下面这样使用 .values 因为我正在传递 pandas DataFrames:

skf = StratifiedKFold(n_splits=3, shuffle=False)
skf.get_n_splits(X, y)

for train_index, test_index, valid_index in skf.split(X.values, y.values):
    print("TRAIN:", train_index, "TEST:", test_index,  "VALID:", valid_index)
    X_train, X_test, X_valid = X.values[train_index], X.values[test_index], X.values[valid_index]
    y_train, y_test, y_valid = y.values[train_index], y.values[test_index], y.values[valid_index]

这失败了:

ValueError: not enough values to unpack (expected 3, got 2).

我通读了所有 sklearn 文档并运行了示例代码，但没有更好地理解如何在 sklearn 之外使用分层 k 折叠拆分交叉验证场景。

编辑:

我也这样试过:

# Create train/test split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.1, stratify=y)

# Create validation split from train split
X_train, X_valid, y_train, y_valid = train_test_split(X_train, y_train, test_size=0.05)

这似乎行得通，尽管我想这样做会扰乱分层。

Answer 1

StratifiedKFold 只能用于将数据集分成两部分。你得到一个错误，因为 split() 方法只会产生一个 train_index 和 test_index 的元组(参见 https://github.com/scikit-learn/scikit-learn/blob/ab93d65/sklearn/model_selection/_split.py#L94 )。

对于此用例，您应该首先将数据拆分为验证数据和其余数据，然后将其余数据再次拆分为测试数据和训练数据，如下所示:

X_rest, X_val, y_rest, y_val = train_test_split(X, y, test_size=0.2, train_size=0.8, stratify=y)
X_train, X_test, y_train, y_test = train_test_split(X_rest, y_rest, test_size=0.25, train_size=0.75, stratify=y_rest)