Pandas数据帧分层分解为训练集、验证集和测试集

Question

以下极度简化的数据框表示包含医疗诊断的更大数据框：

medicalData = pd.DataFrame({'diagnosis':['positive','positive','negative','negative','positive','negative','negative','negative','negative','negative']})
medicalData

    diagnosis
0   positive
1   positive
2   negative
3   negative
4   positive
5   negative
6   negative
7   negative
8   negative
9   negative

问题：对于机器学习，我需要按以下方式将此数据帧随机拆分为三个子帧：

trainingDF, validationDF, testDF = SplitData(medicalData,fractions = [0.6,0.2,0.2])

…其中分割数组指定进入每个子帧的完整数据的分数

• 子帧中的数据需要是互斥的，分割数组（分数）需要求和为一
• 传统上，每个子集的阳性诊断分数需要大致相同。
• Answers to this question 建议使用the pandas sample method或the train_test_split function from sklearn。但这些解决方案似乎都不能很好地推广到n个分裂，也没有一个提供分层分裂

Answer 1

np.array_split

如果您想推广到n拆分，np.array_split是您的朋友（它可以很好地处理数据帧）

fractions = np.array([0.6, 0.2, 0.2])
# shuffle your input
df = df.sample(frac=1) 
# split into 3 parts
train, val, test = np.array_split(
    df, (fractions[:-1].cumsum() * len(df)).astype(int))

---

train_test_split

使用^{}进行分层拆分的多风解决方案

y = df.pop('diagnosis').to_frame()
X = df

X_train, X_test, y_train, y_test = train_test_split(
        X, y,stratify=y, test_size=0.4)

X_test, X_val, y_test, y_val = train_test_split(
        X_test, y_test, stratify=y_test, test_size=0.5)

其中X是功能的数据帧，y是标签的单列数据帧