时间戳

首页 > 解决方案 > 试图将我的数据集拆分为具有代表性的训练集和测试集

python - 试图将我的数据集拆分为具有代表性的训练集和测试集

问题描述

我有一些数据集。我想把它分成一个训练集和测试集。训练集将容纳 2/3 的数据。我希望这两个集合代表整个集合。在我的“类”列中，我有 4 个或 2 个来代表两个类。我希望我的测试集具有相同的 4:2 比例。为了做到这一点，我创建了这段代码：

trainTotal = 455
benTotal = 296
malTotal = 455-296
b = 0
m = 0
tr = 0
i = 0
j = 0

for index, row in data.iterrows():
    if row['Class'] == 2:
        if tr < trainTotal and b < benTotal:
            train.loc[i] = data.iloc[index]
            b = b+1
            tr = tr + 1
            i = i+1
        else:
            test.loc[j] = data.iloc[index]
            j = j+1
    if row['Class'] == 4:
        if tr < trainTotal and m < malTotal:
            train.loc[i] = data.iloc[index]
            tr = tr + 1            
            i = i + 1
            m = m+1
        else:
            test.loc[j] = data.iloc[index]
            j = j + 1

我在我的火车数据框中得到了正确的值总数，但这些案例并没有像我希望的那样表示。它进入了if tr < trainTotal and b < benTotal:太多次。知道问题可能是什么吗？

标签： pythonpandasloops

解决方案

就像迈克尔加德纳说的，train_test_split是你正在寻找的功能。

默认情况下，它会随机拆分，但您可以stratify告诉它您希望训练和测试数据集中的 Class 列具有相同的比率。

它是这样工作的：

from sklearn.model_selection import train_test_split

X_train, X_test, y_train, y_test = train_test_split(
    data,
    target,
    test_size = 0.3,
    stratify=data[['your_column']]
)

推荐阅读