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

Question

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

这是一个例子:

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