python - 如何在 Scikit-Learn 中对训练和测试数据进行分层？

Question

我正在尝试为 Iris 数据集(从 Kaggle 下载)实现分类算法。在物种列中，类别(Iris-setosa、Iris-versicolor、Iris-virginica)按排序顺序排列。如何使用 Scikit-Learn 对训练和测试数据进行分层？

Answer 1

如果你想以 0.3 的测试比率对数据进行洗牌和拆分，你可以使用

sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, shuffle=True)

其中 X 是您的数据，y 是相应的标签，test_size 是应该保留用于测试的数据的百分比，shuffle=True 在之前打乱数据 split

为了确保数据按列均分，可以将其赋予stratify参数。

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.3, 
                                                    shuffle=True, 
                                stratify = X['YOUR_COLUMN_LABEL'])