python - 在 train_test_split 中使用 'stratify' 没有什么区别。它是干什么用的？

Question

我只是在 sklearn 的 train_test_split 函数中尝试“分层”参数。我的数据集不平衡，以下是类别的比例:

0级:8,902 1级:1,605

第 1 类占数据集的 15%。

这是不使用分层的默认分割:

x_train, x_test, y_train, y_test = train_test_split(df['image'], df['class'], test_size=0.2,random_state=5)

Training set balance:
0    7,116
1    1,289

Test set balance:
0    1,786
1     316

下面我使用分层:

x_train, x_test, y_train, y_test = train_test_split(df['image'], df['class'], test_size=0.2,random_state=5,stratify=df['class'])

Training set balance:
0    7121
1    1284

Test set balance:
0    1781
1     321

两者的比例大致相同:第 1 类为 18%。添加“分层”没有任何作用。

所以这让我有点困惑。我做错了什么吗？

谢谢

Answer 1

添加stratify将保证1的比例与原始数据相同。

计算1的比例:

原文:

Total:  print(1605/(1605+8902)) = 0.1527553059864852

没有分层:

Train:  print(1289/(1289+7116)) = 0.1533610945865556
Test:   print(316/(316+1786)) = 0.15033301617507136

正如你所看到的，1的比例与原始数据不一样，当你再次采样时，比例可能会不同! (因为是随机抽样所以很相似)

分层:

Train:  print(1284/(1284+7121)) = 0.15276621058893516
Test:   print(321/(321+1781)) = 0.1527117031398668

与原始数据一样，即使再次采样，比例也不会改变。那么分层就发挥了它的作用，不是吗？