sklearn中 KFold 和 StratifiedKFold 差别

KFold划分数据集:根据n_split直接进行顺序划分,不考虑数据label分布
StratifiedKFold划分数据集:划分后的训练集和验证集中类别分布尽量和原数据集一样

验证:

from sklearn.model_selection import KFold
from sklearn.model_selection import StratifiedKFold
import numpy as np

X = np.array([[10, 1], [20, 2], [30, 3], [40, 4], [50, 5], [60, 6], [70, 7], [80, 8], [90, 9], [100, 10], [90, 9], [100, 10]])

# 两个类别:1:1
Y = np.array([1,1,1,1,1,1,2,2,2,2,2,2])

print("Start Testing KFold...")
# KFold划分数据集的原理:根据n_split直接进行顺序划分
kfolds = KFold(n_splits=3, shuffle=False)
for (trn_idx, val_idx) in kfolds.split(X, Y):
    print((trn_idx, val_idx))
    print((len(trn_idx), len(val_idx)))


print('\n' + "Start Testing StratifiedKFold...")
# StratifiedKFold: 抽样后的训练集和验证集的样本分类比例和原有的数据集尽量是一样的
stratifiedKFolds = StratifiedKFold(n_splits=3, shuffle=False)
for (trn_idx, val_idx) in stratifiedKFolds.split(X, Y):
    print((trn_idx, val_idx))
    print((len(trn_idx), len(val_idx)))

out:

sklearn中 KFold 和 StratifiedKFold 差别_第1张图片
note:
另外使用StratifiedKFold时建议使用shuffle=True,个人实践发现设置为false后训练集和验证集acc会大幅下降(从0.8降到0.4)比较离谱;但是测试集却变化不大;目前还没找着原因…

你可能感兴趣的:(pytorch,KFold,StratifiedKFold)