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

Question

我正在尝试实现虹膜数据集的分类算法（从 Kaggle 下载）。在“物种”列中，类别（山鸢尾、杂色鸢尾、维吉尼亚鸢尾）按排序顺序排列。如何使用 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在分割之前对数据进行洗牌

为了确保数据按列平均分割，您可以将其赋予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'])