Sklearn 的课堂“StratifiedShuffleSplit”

Question

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

下面的代码来自 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]

出来：

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 的拆分。比如第一次split对原始数据打乱，样本5,2,3被选为训练集，这也是group_label的分层抽样；在第二次拆分中，数据再次洗牌，样本5,1,4被选为训练集；等等。

Answer 2

split.split() 函数返回训练样本和测试样本的索引。它会查看指定的交叉验证次数，并在每次训练和测试样本索引时返回，使用哪个训练和测试数据集可以通过过滤整个数据集来创建。