python – 在训练,验证和测试集中分层分离pandas数据帧

以下极其简化的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])

在分割阵列指定进入每个子帧的完整数据的分数的情况下,子帧中的数据需要互斥,并且分割阵列需要总和为1.
另外,每个子集中阳性诊断的比例需要大致相同.

Answers to this question 建议使用the pandas sample method或the train_test_split function from sklearn.但是这些解决方案似乎都不能很好地推广到n个分裂,并且没有一个提供分层分割.

解决方法:

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)

其中X是要素的DataFrame,y是标签的单柱数据框架.