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

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

我使用以下代码来构建模型的测试和训练数据集。

from sklearn.cross_validation import train_test_split X = Actual_DataFramey = 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’的百分比各是多少。

因此，需要一个函数，该函数接受2个输入来创建训练数据：-训练数据的总记录数，-属于类别'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

有更节省内存或运行更快的方法，我没有对这段代码进行任何测试，但它应该能工作。

至少，你应该能理解背后的思路。