sklearn的train_test_split中的“分层”参数无法正常工作？

Question

我stratify在train_test_split()scikit-learn函数中的参数有问题。这是一个具有相同问题的虚拟示例，该问题随机出现在我的数据上：

from sklearn.model_selection import train_test_split
a = [1, 0, 0, 0, 0, 0, 0, 1]
train_test_split(a, stratify=a, random_state=42)

返回：

[[1, 0, 0, 0, 0, 1], [0, 0]]

它不应该在测试子集中也选择“ 1”吗？从我多么希望train_test_split()有stratify工作就应该返回类似：

[[1, 0, 0, 0, 0, 0], [0, 1]]

对于某些值，会发生这种情况random_state，而对于其他值，它会正常工作；但是我每次必须分析数据时都无法搜索它的“正确”值。

我有python 2.7和scikit-learn 0.18。

Answer 1

这个问题是在8个月前提出的，但是我想一个答案将来可能仍然对读者有帮助。

使用stratify参数时，train_test_split实际上是依靠StratifiedShuffleSplit函数进行拆分的。正如您在文档中所看到的，StratifiedShuffleSplit确实旨在通过保留每个类别的样本百分比来进行拆分，这与您期望的一样。

问题是，在您的示例中，25％（8个样本中的2个）为1s，但是样本大小不足以使您看到此比例反映在测试集上。您在这里有两个选择：

A.使用选项增大测试集的大小，该选项test_size默认为0.25，例如0.5。在这种情况下，一半样本将成为测试集，您会看到其中25％（即四分之一）是1。

>>> a = [1, 0, 0, 0, 0, 0, 0, 1]
>>> train_test_split(a, stratify=a, random_state=42, test_size=0.5)
[[1, 0, 0, 0], [0, 0, 1, 0]]

B.保持test_size其默认值并增加集合的大小，a以便其样本的25％至少等于4个元素。a16个或更多样本中的一个将为您完成此任务。

>>> a = [1, 0, 0, 0, 0, 0, 0, 1, 1, 0, 0, 0, 0, 0, 0, 1]
>>> train_test_split(a, stratify=a, random_state=42)
[[0, 0, 1, 0, 0, 1, 0, 1, 0, 0, 0, 0], [0, 0, 1, 0]]

希望能有所帮助。