Sklearn 的课程“StratifiedShuffleSplit”

Question

我StratifiedShuffleSplit对 Sklearn类的工作原理有点困惑。

下面的代码来自 Géron 的书“Hands On Machine Learning”，第 2 章，他在那里进行了分层抽样。

from sklearn.model_selection import StratifiedShuffleSplit

split = StratifiedShuffleSplit(n_splits=1, test_size=0.2, random_state=42)
for train_index, test_index in split.split(housing, housing["income_cat"]):
    strat_train_set = housing.loc[train_index]
    strat_test_set = housing.loc[test_index]

特别是，在做什么split.split？

谢谢！

Answer 1

由于您没有提供数据集，我使用 sklearn 示例来回答这个问题。

准备数据集

# generate data
import numpy as np
from sklearn.model_selection import StratifiedShuffleSplit
data = np.array([[1, 2], [3, 4], [1, 2], [3, 4], [1, 2], [3, 4]])
group_label = np.array([0, 0, 0, 1, 1, 1])

这会生成一个数据集data，其中有 6 个观察和 2 个变量。group_label有 2 个值，均值group 0和group 1。在这种情况下，group 0包含 3 个样本，同样是group 1。一般来说，组大小不需要相同。

创建StratifiedShuffleSplit对象实例

sss = StratifiedShuffleSplit(n_splits=5, test_size=0.5, random_state=0)
sss.get_n_splits(data, group_label)

出去：

5

在这一步，你可以创建一个 的实例StratifiedShuffleSplit，你可以告诉函数如何拆分（在random_state = 0，拆分数据5 times，每次50%数据都会拆分为test set）。但是，它只会在您在下一步中调用它时拆分数据。

调用实例，拆分数据。

# the instance is actually a generater
type(sss.split(data, group_label))

# split data
for train_index, test_index in sss.split(data, group_label):
     print("n_split",,"TRAIN:", train_index, "TEST:", test_index)
     X_train, X_test = X[train_index], X[test_index]
     y_train, y_test = y[train_index], y[test_index]
type(sss.split(data, group_label))

出去：

TRAIN: [5 2 3] TEST: [4 1 0]
TRAIN: [5 1 4] TEST: [0 2 3]
TRAIN: [5 0 2] TEST: [4 3 1]
TRAIN: [4 1 0] TEST: [2 3 5]
TRAIN: [0 5 1] TEST: [3 4 2]

在这一步中，spliter您在上一步中定义的将生成 5 个data一一拆分。例如，在第一次拆分中，将原始数据打乱，并5,2,3选择样本作为训练集；在第二次拆分中，数据再次打乱，样本5,1,4被选为训练集；等等..