"Stratify" sklearn train_test_split中的参数无法正常工作？

Question

我对scikit-learn的train_test_split()函数中的stratify参数有一个问题。这是一个虚拟的例子，与我的数据随机出现的问题相同:

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个)是15，但样本量不够大，您无法在测试集中看到这个比例。这里有两个选项:

。使用选项test_size增加测试集的大小，默认为0.25，例如0.5。在这种情况下，一半的样本将成为测试集，并且您将看到其中25%(即1/4)为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个元素。包含16个或更多样本的a将为您完成此操作。

>>> 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]]

希望对你有帮助。