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

Posted 2023-03-12

【中文标题】使用 StratifiedKFold 创建训练/测试/验证拆分

【英文标题】：Creating train/test/val split with StratifiedKFold

【发布时间】：2017-12-26 13:51:42

【问题描述】：

我正在尝试使用 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)

这似乎可行，尽管我想我这样做会弄乱分层。

【问题讨论】：

您的问题到底是什么？这与您的预期有何不同？

感谢您的回复@RyanStout，我的错误 - 已更新错误消息。

【参考方案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)

【讨论】：

我不确定在这里使用stratify='column'，但是当我在我的数据上运行你的代码时，我得到：TypeError: Singleton array array('column', dtype='<U6') cannot be considered a valid collection.

我也有同样的情况@tw0000

【参考方案2】：

在stratify参数中，传递目标进行分层。首先，通知完整的目标数组（在我的例子中是y）。然后，在下一次拆分中，通知被拆分的目标（在我的情况下为y_train）：

X = df.iloc[:,:-1].values
y = df.iloc[:,-1].values

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42, stratify=y)

X_train, X_val, y_train, y_val = train_test_split(X_train, y_train, test_size=0.2, random_state=42, stratify=y_train)

【参考方案3】：

我不确定这个问题是关于 KFold 还是只是分层拆分，但我用交叉验证集为 StratifiedKFold 编写了这个快速包装器。

from sklearn.model_selection import StratifiedKFold, train_test_split

class StratifiedKFold3(StratifiedKFold):

    def split(self, X, y, groups=None):
        s = super().split(X, y, groups)
        for train_indxs, test_indxs in s:
            y_train = y[train_indxs]
            train_indxs, cv_indxs = train_test_split(train_indxs,stratify=y_train, test_size=(1 / (self.n_splits - 1)))
            yield train_indxs, cv_indxs, test_indxs

可以这样使用：

X = np.random.rand(100)
y = np.random.choice([0,1],100)
g = KFold3(10).split(X,y)
train, cv, test = next(g)
train.shape, cv.shape, test.shape
>> ((80,), (10,), (10,))