我在尝试使用 sklearn 的 train_test_split 函数中的 ‘stratified’ 参数。我的数据集是不平衡的，以下是各类别的比例：

类别 0：8,902 类别 1：1,605

类别 1 占数据集的 15%。

以下是默认分割情况，未使用 stratify：

x_train, x_test, y_train, y_test = train_test_split(df['image'], df['class'], test_size=0.2,random_state=5)训练集平衡情况：0    7,1161    1,289测试集平衡情况：0    1,7861     316

下面我使用了 stratify：

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'])训练集平衡情况：0    71211    1284测试集平衡情况：0    17811     321

在两种情况下，比例大致相同：类别 1 占 18%。添加 ‘stratify’ 并没有起到作用。

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

谢谢

回答：

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

计算类别 1 的比例：

原始数据：

总数： print(1605/(1605+8902)) = 0.1527553059864852

未使用 stratify：

训练集： print(1289/(1289+7116)) = 0.1533610945865556测试集： print(316/(316+1786)) = 0.15033301617507136

如您所见，类别 1 的比例与原始数据不同，并且在再次抽样时，比例可能会有所不同！（它相似是因为这是随机抽样）

使用 stratify：

训练集： print(1284/(1284+7121)) = 0.15276621058893516测试集： print(321/(321+1781)) = 0.1527117031398668

与原始数据相同，即使再次抽样，比例也不会改变。所以 stratify 确实起到了作用，不是吗？