python - 来自 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 个月前提出的，但我想答案可能对 future 的读者有所帮助。

当使用stratify 参数时，train_test_split 实际上依赖于StratifiedShuffleSplit 函数来进行拆分。正如您在 documentation 中看到的那样, StratifiedShuffleSplit 确实旨在通过保留每个类的样本百分比来进行拆分，正如您所期望的那样。

问题是，在您的示例中，25%(8 个样本中的 2 个)为 1，但样本量不够大，您无法在测试集上看到这一比例。您在这里有两个选择:

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 个元素。 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]]

希望对您有所帮助。