Question

我需要将数据拆分为训练集（75％）和测试集（25％）。我目前使用以下代码执行此操作：

X, Xt, userInfo, userInfo_train = sklearn.cross_validation.train_test_split(X, userInfo)

但是，我想对训练数据集进行分层。我怎么做？我一直在调查StratifiedKFold方法，但不允许我指定75％/ 25％的分割，只对训练数据集进行分层。

Answer 1

[更新0.17]

请参阅sklearn.model_selection.train_test_split的文档：

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y,
                                                    stratify=y, 
                                                    test_size=0.25)

[/更新0.17]

有拉取请求here。但你可以做train, test = next(iter(StratifiedKFold(...))) 如果你愿意，可以使用火车和测试指数。

Answer 2

TL; DR：将StratifiedShuffleSplit与test_size=0.25

一起使用

Scikit-learn为分层拆分提供了两个模块：

StratifiedKFold：此模块可用作直接k折叠交叉验证运算符：因为它将设置n_folds训练/测试集，以便类在两者中均衡。

下面是一些代码（直接来自上面的文档）

>>> skf = cross_validation.StratifiedKFold(y, n_folds=2) #2-fold cross validation
>>> len(skf)
2
>>> for train_index, test_index in skf:
...    print("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]
...    #fit and predict with X_train/test. Use accuracy metrics to check validation performance

StratifiedShuffleSplit：此模块创建一个具有同等平衡（分层）类的训练/测试集。基本上这就是你想要的n_iter=1。您可以在此处提及与train_test_split

代码：

>>> sss = StratifiedShuffleSplit(y, n_iter=1, test_size=0.5, random_state=0)
>>> len(sss)
1
>>> for train_index, test_index in sss:
...    print("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]
>>> # fit and predict with your classifier using the above X/y train/test

Answer 3

以下是连续/回归数据的示例（直到this issue on GitHub被解析）。

# Your bins need to be appropriate for your output values
# e.g. 0 to 50 with 25 bins
bins     = np.linspace(0, 50, 25)
y_binned = np.digitize(y_full, bins)
X_train, X_test, y_train, y_test = train_test_split(X, y, stratify=y_binned)

Answer 4

除了@Andreas Mueller接受的回答之外，只想将其添加为上面提到的@tangy：

StratifiedShuffleSplit最接近train_test_split（stratify = y）增加了以下功能：

默认情况下分层
通过指定 n_splits ，它会重复拆分数据

Answer 5

您可以简单地使用Scikit Learn中可用的train_test_split()方法来做到这一点：

from sklearn.model_selection import train_test_split 
train, test = train_test_split(X, test_size=0.25, stratify=X['YOUR_COLUMN_LABEL'])

我还准备了一个简短的GitHub Gist，其中显示了stratify选项的工作原理：

https://gist.github.com/SHi-ON/63839f3a3647051a180cb03af0f7d0d9

Answer 6

StratifiedShuffleSplit 是在我们选择应该在我们即将生成的所有小数据集中均匀表示的列之后完成的。 '折叠是通过保留每个类的样本百分比来制作的。'

假设我们有一个带有“季节”列的数据集“数据”，并且我们希望得到“季节”的均匀表示，那么它看起来像这样：

from sklearn.model_selection import StratifiedShuffleSplit
sss=StratifiedShuffleSplit(n_splits=1,test_size=0.25,random_state=0)

for train_index, test_index in sss.split(data, data["season"]):
    sss_train = data.iloc[train_index]
    sss_test = data.iloc[test_index]

Answer 7

#train_size is 1 - tst_size - vld_size
tst_size=0.15
vld_size=0.15

X_train_test, X_valid, y_train_test, y_valid = train_test_split(df.drop(y, axis=1), df.y, test_size = vld_size, random_state=13903) 

X_train_test_V=pd.DataFrame(X_train_test)
X_valid=pd.DataFrame(X_valid)

X_train, X_test, y_train, y_test = train_test_split(X_train_test, y_train_test, test_size=tst_size, random_state=13903)

Answer 8

从上方将@tangy答案更新为scikit-learn的当前版本：0.23.2（StratifiedShuffleSplit documentation）。

from sklearn.model_selection import StratifiedShuffleSplit

n_splits = 1  # We only want a single split in this case
sss = StratifiedShuffleSplit(n_splits=n_splits, test_size=0.25, random_state=0)

for train_index, test_index in sss.split(X, y):
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

Answer 9

因此，最好将数据集拆分为训练集和测试集，以保持每个类中的示例比例与在原始数据集中观察到的比例相同。

这称为分层训练-测试拆分。

我们可以通过将“stratify”参数设置为原始数据集的 y 分量来实现这一点。这将由 train_test_split() 函数使用，以确保训练集和测试集都具有所提供的“y”数组中存在的每个类中的示例比例。

在scikit-learn中分层训练/测试分裂

9 个答案: