KFold StratifiedKFold GroupKFold的应用场景及示例代码

 

1. KFold

普通的k折交叉验证

from sklearn.model_selection import KFold

kf = KFold(n_splits=5, shuffle=True, random_state=8)

for train_index, test_index in kf.split(X):
    print("TRAIN:", train_index, "TEST:", test_index)
    X_train, X_test = X[train_index], X[test_index]
    y_train, y_test = y[train_index], y[test_index]

 

2. StratifiedKFold

平衡label的k折交叉验证。在类别label不平衡时使用，如0、1分类时，0的样本量和1的样本量差异悬殊。

from sklearn.model_selection import StratifiedKFold

skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=3)

for i, (train_index, valid_index) in enumerate(skf.split(X_train, y_train)):
    X_trn, X_val = X_train.iloc[train_index], X_train.iloc[valid_index]
    y_trn, y_val = y_train.iloc[train_index], y_train.iloc[valid_index]

 

3. GroupKFold 

比如df数据集中，有一列用户的编号user_id，而其中每个user_id可能有多个样本，我们在使用k折交叉验证的时候希望对user_id进行拆分（类似groupby），将一个user_id的所有样本同时放在train中或者vaild中。

from sklearn.model_selection import GroupKFold

group_kfold = GroupKFold(n_splits=5)

for i, (train_index, valid_index) in enumerate(group_kfold.split(X_train, y_train, df['user_id'])):

        X_trn, X_val = X_train.iloc[train_index], X_train.iloc[valid_index]
        y_trn, y_val = y_train.iloc[train_index], y_train.iloc[valid_index]