python-3.x - Python - 自定义采样以获取训练和测试数据

Question

我有一个高度不平衡的数据集。

我的数据集包含 1450 条记录，我的输出是二进制 0 和 1。输出 0 有 1200 条记录，1 有 250 条记录。

我正在使用这段代码来构建模型的测试和训练数据集。

from sklearn.cross_validation import train_test_split 
X = Actual_DataFrame
y = Actual_DataFrame.pop('Attrition')
X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.20, random_state=42, stratify=y)

但是我想要的是一种通过函数的方法，在该函数中我想指定训练记录的数量以及其中有多少百分比需要来自“0”类以及其中有多少百分比需要来自类“1”。

So, a function which takes 2 Inputs are needed for creating the training_data:-
Total Number of Records for Training Data,
Number of Records that belongs to Class '1'

这对于解决有偏差的采样数据集问题将有巨大帮助。

Answer 1

您可以简单地编写一个与 sklearn 中的 train_test_split 非常相似的函数。这个想法是，根据输入参数 train_size 和 pos_class_size，您可以计算出需要多少正类样本和负类样本。

def custom_split(X, y, train_size, pos_class_size, random_state=42):
    neg_class_size = train_size = pos_class_size
    pos_df = X[y == 1]
    neg_df = X[y == 0]

    pos_train = pos_df.sample(pos_class_size)
    pos_test = pos_df[~pos_df.index.isin(pos_train.index)]

    neg_train = neg_df.sample(neg_class_size)
    neg_test = neg_df[~neg_df.index.isin(neg_train.index)]

    X_train = pd.concat([pos_train,neg_train], axis=1)
    X_test = pd.concat([pos_test,neg_test], axis=1)

    y_train = y[X_train.index]
    y_test = y[X_test.index]

    return X_train, X_test, y_train, y_test

有些方法内存效率高或运行速度更快，我没有对此代码进行任何测试，但它应该可以工作。

至少，你应该能够理解背后的想法。