样本类别不均匀使用StratifiedKFold分层

前言

本文采用代码的方式理解KFold和StratifiedKFold的不同与作用。

简述KFold和StratifiedKFold

KFold和StratifiedKFold都是一个交叉验证器，它将数据集分成k个折叠。而分层是为了确保数据集的每一个部分都有相同比例的观察值和给定的标签。因此，这意味着StratifiedKFold是KFold的改进版本。
假设有一个具有16个数据点和不平衡类分布的数据集。在数据集中，12个数据点属于A类，其余（即4个）属于B类。B类与A类的比率为1/3。如果我们使用StratifiedKFold和集合k=4，那么训练集将包括来自A类的3个数据点和来自B类的9个数据点，而测试集包括来自A级的3个和来自B级的1个数据点。

代码展示

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],
    [61,62,63,64],
    [71,72,73,74],
    [71,72,73,74],
    [71,72,73,74],
    [71,72,73,74],
    [71,72,73,74],
    [71,72,73,74],
    [71,72,73,74],
])

y1 = np.array([1,1,2,2,3,3,4,4,5,5,6,6,7,7,8,8]) # 1：1：1：1：1：1:1:1
y2 = np.array([1,1,1,1,2,2,2,2,2,2,2,2,2,2,2,2]) # 1:3
sfolder = StratifiedKFold(n_splits=2,random_state=2,shuffle=True)
floder = KFold(n_splits=2,random_state=2,shuffle=True)

# 按照1：1：1：1：1：1:1:1划分
for train, test in sfolder.split(X,y1):
    print('Train: %s | test: %s' % (train, test))
print("StratifiedKFold done")

for train, test in floder.split(X,y1):
    print('Train: %s | test: %s' % (train, test))
print("KFold done")

# 按照 1:3
for train, test in sfolder.split(X,y2):
    print('Train: %s | test: %s' % (train, test))
print("StratifiedKFold done")

for train, test in floder.split(X,y2):
    print('Train: %s | test: %s' % (train, test))
print("KFold done")

结果分析

可以从结果看到不管是从样本均匀的或者样本不均匀的情况下，StratifiedKFold总能按照原始数据的比例精确划分，原始数据有16个，n_splits=2，所以会划分成两组，y1的类别比例是1:1:1:1:1:1:1:1，用StratifiedKFold划分，每组是八个，比例也会是1:1:1:1:1:1:1:1，而用KFold则不是。当样本类别不均匀时如y2类别比例是1:3，用StratifiedKFold划分得到的两组比例仍是1:3。

总结
在处理具有不平衡类分布的分类任务时，我们应该更喜欢StratifiedKFold而不是KFold。