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

Posted 2023-03-12

【中文标题】试图将我的数据框拆分为具有代表性的训练集和测试集

【英文标题】：Trying to split my datafame into a representative train and test set

【发布时间】：2020-02-07 11:00:39

【问题描述】：

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

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

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

【问题讨论】：

为什么不使用 sklearn train_test_split？ ***.com/questions/29438265/…

输入超过296次吗？

我想如果我使用 train_test_split 它只会按照我指定的方式拆分数据。我担心数据不会像我希望的那样代表类列

是的，我不知道为什么会@cosmic_inquiry

@newwebdev22 阅读我提供的链接。有很多选项可以对您的数据进行分层。

【参考方案1】：

就像 Michael Gardner 说的，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']]
)