使用 sklearn.model_selection 拆分不平衡的数据集答案

【问题标题】：Using sklearn.model_selection to split unbalanced dataset使用 sklearn.model_selection 拆分不平衡的数据集
【发布时间】：2019-09-25 04:56:45
【问题描述】：

我正在使用以下代码将我的数据集拆分为训练/验证/测试集。

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = 
        train_test_split(X_data, y_data, test_size=0.3, random_state=42)

X_test, X_val, y_test, y_val = 
        train_test_split(X_test, y_test, test_size=0.5, random_state=42)

问题是我的数据集真的不平衡。例如，有些类有 500 个样本，而有些类有 70 个。这种拆分方法在这种情况下准确吗？抽样是随机的还是 sklearn 使用 seome 方法来保持数据在所有集合中的分布相同？

【问题讨论】：

为什么要将测试数据拆分两次？你会得到这样的 15% 的测试数据大小吗？
@Scotty1 目的是获取训练集（70%）、验证集（15%）和测试集（15%）
如果您可以使用“不平衡数据集”更新问题，那就太好了。将来会对其他人有所帮助。
@Aditya Kansal 做到了
回答没有帮助？

标签： python machine-learning scikit-learn dataset

【解决方案1】：

您应该使用stratify 选项（请参阅docs）：

X_train, X_test, y_train, y_test = 
        train_test_split(X_data, y_data, test_size=0.3, random_state=42, stratify=y_data)

X_test, X_val, y_test, y_val = 
        train_test_split(X_test, y_test, test_size=0.5, random_state=42, stratify=y_test)

【讨论】：