Question

我在scikit-learn的stratify函数中遇到train_test_split()参数问题。这是一个虚拟示例，在我的数据中出现了相同的问题：

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”吗？从{I} 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个）是1，但样本大小不足以让您看到此比例反映在测试集上。你有两个选择：

A. 使用选项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]]

希望有所帮助。

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

1 个答案: