StratifiedKfold和KFold的区别

@@机器学习小白

StratidfiedKfold和KFold交叉验证，都是将训练集分成K份，其中K-1份的训练集，1份的验证集。

不过特别的是StratifiedKfold将验证集的正负样本比例，保持和原始数据的正负样本比例相同

import numpy as np
from sklearn.model_selection import StratifiedKFold,KFold
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,0,0,0,1,0,0,0])

folder = KFold(n_splits=4,random_state=0,shuffle=True)
sfolder = StratifiedKFold(n_splits=4,random_state=0,shuffle=True)
for train,test in sfolder.split(y,y):
    print((train,test))
print('----------')
for train,test in folder.split(y,y):
    print(train,test)

##output
(array([0, 1, 2, 4, 5, 6]), array([3, 7]))
(array([0, 1, 3, 4, 6, 7]), array([2, 5]))
(array([0, 2, 3, 5, 6, 7]), array([1, 4]))
(array([1, 2, 3, 4, 5, 7]), array([0, 6]))
----------
[0 1 3 4 5 7] [2 6]
[0 2 3 4 5 6] [1 7]
[1 2 4 5 6 7] [0 3]
[0 1 2 3 6 7] [4 5]

从上述结果可以看到，StratifiedKfold划分出来的验证集里的正负样本比例和原始数据集的一样，而KFold没有