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

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函数进行拆分。正如您在documentation 中看到的那样，StratifiedShuffleSplit 确实旨在通过保留每个类的样本百分比来进行拆分，正如您所期望的那样。

问题是，在您的示例中，25%（8 个样本中的 2 个）是 1s，但样本量不足以让您在测试集上看到这一比例。您有两种选择：

A. 使用选项test_size（默认为 0.25）增加测试集的大小，例如 0.5。在这种情况下，一半的样本将成为您的测试集，您会看到其中 25%（即 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]]

希望对您有所帮助。