试图将我的数据分成代表性的训练和测试

Question

我有一些数据集。我想把它分成一列火车和一组测试。列车组将保存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:的次数太多了。你知道问题是什么吗

Answer 1

正如Michael Gardner所说，train_test_split是您正在寻找的函数

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

它的工作原理如下：

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']]
)