Sklearn中的类"StratifiedShuffleSplit"

Question

我对Sklearn库中的StratifiedShuffleSplit类如何工作感到有些困惑。

下面的代码来自Géron的书《动手学机器学习》，第二章中，他进行了分层抽样。

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 中做了什么？

谢谢！

- Rafael Higa

---

这个回答解决了你的问题吗？sklearn中StratifiedKFold和StratifiedShuffleSplit的区别 - PV8

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次，每次将50％的数据分割到测试集）。但是，只有当您在下一步中调用它时才会分割数据。

调用该实例并分割数据。

# 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”进行分割。例如，在第一个分割中，原始数据将被打乱并且样本“5,2,3”被选择为训练集，这也是根据“group_label”进行分层抽样的；在第二个分割中，数据再次被打乱，并选择了样本“5,1,4”作为训练集；以此类推。

- Travis

Answer 2

split.split()函数返回训练样本和测试样本的索引。它将查找指定的交叉验证次数，并每次返回用于通过筛选整个数据集创建训练和测试数据集的训练和测试样本索引。

- Sunnysinh Solanki