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是标签的单列数据帧。