Question

我正在尝试以分层方式拆分数据。我认为train_test_split中的sklearn在不平衡的数据集上无法按预期方式工作。

这里是一个例子：

from sklearn.model_selection import train_test_split
from collections import Counter

y = [7]*2 + [1]*100 + [2]*3 + [3]*3 + [4]*6 + [5]*100 + [6]*2 + [8]*2
xtrain, xtest = train_test_split(
        y, stratify=y, test_size=0.2, 
        shuffle=True, random_state=42
    )

print(Counter(xtrain))
#Counter({1: 80, 5: 80, 2: 2, 4: 5, 8: 2, 3: 2, 7: 2, 6: 1})

print(Counter(xtest))
#Counter({5: 20, 1: 20, 6: 1, 2: 1, 3: 1, 4: 1})

第1类和第5类按预期方式分配。但是，代表性不足的类分布不均。类6,2,3,4分布良好，但第8类和第7类却分布不好。

Answer 1

我实际上认为它确实有效。在某些类别中，样本太少了，以致于无法在统计上准确。

我尝试了以下

y = [7]*200 + [1]*10000 + [2]*300 + [3]*300 + [4]*600 + [5]*10000 + [6]*200 + [8]*200

输出正确：

>>> Counter({1: 8000, 5: 8000, 4: 480, 2: 240, 3: 240, 6: 160, 7: 160, 8: 160})
Counter({1: 2000, 5: 2000, 4: 120, 2: 60, 3: 60, 6: 40, 7: 40, 8: 40})

最好的问候 -ga97dil

Answer 2

我认为您没有丢失任何东西。 Stratify尝试在拆分期间保持每个类的表示形式。唯一的问题是样本大小不足以让您看到测试大小的完整反映。

您可以通过增加8和7的样本量来解决此问题。

train_test_split在分层数据上无法按预期工作

2 个答案: