python - 将 pandas 数据帧分层拆分为训练、验证和测试集

Question

以下极其简化的 DataFrame 代表了一个更大的包含医疗诊断的 DataFrame：

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

...其中拆分数组指定进入每个子帧的完整数据的一部分。

• 子帧中的数据需要互斥，并且拆分数组（分数）需要求和为一。
• 此外，每个子集中的阳性诊断比例需要大致相同。
• 此问题的答案建议使用pandas 示例方法或sklearn 中的 train_test_split 函数。但是这些解决方案似乎都不能很好地推广到 n 个拆分，也没有一个提供分层拆分。

标签： pythonpandasdataframemachine-learningsampling

Answer 1

np.array_split

如果您想推广到n拆分，np.array_split是您的朋友（它可以很好地与 DataFrames 一起使用）。

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

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)

WhereX是您的功能的 DataFrame，并且y是您的标签的单列 DataFrame。