概述:
StratifiedKFold用法类似Kfold,但是他是分层采样,确保训练集,测试集中各类别样本的比例与原始数据集中相同。
import numpy as np from sklearn.model_selection import KFold,StratifiedKFold X=np.array([ [1,2,3,4], [11,12,13,14], [21,22,23,24], [31,32,33,34], [41,42,43,44], [51,52,53,54], [61,62,63,64], [71,72,73,74] ]) y=np.array([1,1,0,0,1,1,0,0]) sfolder=StratifiedKFold(n_splits=4,random_state=0,shuffle=False) for train, test in sfolder.split(X,y): print(train, test) OUT: [1 3 4 5 6 7] [0 2] [0 2 4 5 6 7] [1 3] [0 1 2 3 5 7] [4 6] [0 1 2 3 4 6] [5 7] floder = KFold(n_splits=4,random_state=0,shuffle=False) for train, test in floder.split(X,y): print('Train: %s | test: %s' % (train, test)) OUT: Train: [2 3 4 5 6 7] | test: [0 1] Train: [0 1 4 5 6 7] | test: [2 3] Train: [0 1 2 3 6 7] | test: [4 5] Train: [0 1 2 3 4 5] | test: [6 7]
注意返回的仅仅是索引号,可以看到上图中StratifiedKFold 分层采样交叉切分,确保训练集,测试集中各类别样本的比例与原始数据集中相同。比如原数据中,0,1两类比例是1:1,通过观察StratifiedKFold切分的每个测试集可以发现,0,1两类的占比也为1:1,这就是分层采样。确定了测试集后,测试集的补集就是训练集。同时应注意到每个测试集均是互斥的。