应用分层K折交叉验证后,如何将数据分为测试和训练?

时间:2020-08-18 02:41:29

标签: python pandas scikit-learn k-fold

我已经使用以下代码将列分配给它们的特定k折叠:

from sklearn.model_selection import StratifiedKFold, train_test_split

# Stratified K-fold cross-validation 
df['kfold'] = -1
df = df.sample(frac=1).reset_index(drop=True)
y = df.quality
kf = StratifiedKFold(n_splits=5)

for f, (t_,v_) in enumerate(kf.split(X=df, y=y)):
  df.loc[v_, 'kfold'] = f

现在数据框符合预期:


        alcohol  volatile acidity   sulphates citric acid   quality kfold
1499    10.9            0.36          0.73        0.39          6   4
1500    9.5             0.65          0.55        0.10          5   4
1501    13.4            0.44          0.66        0.68          6   4
1502    9.6             0.59          0.67        0.24          5   4
1503    13.0            0.53          0.77        0.79          5   4

但是我如何将其拆分为训练和测试拆分?

1 个答案:

答案 0 :(得分:1)

StratifiedKFold将数据框分成多个折叠,并返回训练/测试索引。 每折将有一部分用于测试(大小为len(data)/ n),其余部分用于训练。

在for循环中,您可以按如下方式访问训练和测试集:

for f, (t_,v_) in enumerate(kf.split(X=df, y=y)):
  df_train = df.loc[t_]
  df_test = df.loc[v_]

您会看到kfold列中添加了标签测试数据。其余数据应用于此折叠的训练。也就是说,对于kfold == 1,训练数据是所有其他数据(kfold != 1)。