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。

- Hantaa

---

如果您尝试使用 stratify=np.unique(a) 会怎样呢？ - MMF

很遗憾，它无法工作，因为传递给“分层”的列表必须与要拆分的列表长度相同。 - Hantaa

文档中没有任何地方声明即使在非常小的子集中也会有所有类。如果您将仅1添加到列表中，则在测试拆分中获得类1。我认为它应该是与训练拆分中的类1相同的部分。例如，如果您删除“stratify”，则会获得列表的尾部而不是带有洗牌类的列表。 - sergzach

Answer 1

这个问题是8个月前提出的，但我认为答案可能仍然有助于未来的读者。

在使用参数stratify时，train_test_split实际上依赖于StratifiedShuffleSplit函数来进行拆分。正如您在文档中所见，StratifiedShuffleSplit旨在通过保留每个类别样本的百分比来进行拆分，就像您期望的那样。

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

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

希望这有所帮助。

- DanielP

---

谢谢你的回答！非常有帮助！！ - msarafzadeh