scikit-learn中的分层训练/测试拆分

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]

有一个拉要求在这里。但是，您可以根据需要简单地进行train, test = next(iter(StratifiedKFold(...))) 训练和测试索引的使用。