StratifiedKFold#

class sklearn.model_selection.StratifiedKFold(n_splits=5, *, shuffle=False, random_state=None)[source]#

按类别分层的 K-折交叉验证器。

提供训练/测试索引，用于将数据拆分为训练/测试集。

此交叉验证对象是 KFold 的一种变体，它返回分层折叠。折叠是通过在二元或多类别分类设置中，保留 y 中每个类别的样本百分比来创建的。

在用户指南中阅读更多内容。

有关交叉验证行为的可视化以及常见 scikit-learn 分割方法的比较，请参阅在 scikit-learn 中可视化交叉验证行为

注意

基于类别标签的分层解决的是一个工程问题而非统计问题。有关更多详细信息，请参阅基于类别标签分层的交叉验证迭代器。

参数:

n_splitsint, 默认为 5: 折叠数量。必须至少为 2。

0.22 版本中的更改: n_splits 的默认值从 3 更改为 5。
shufflebool, 默认为 False: 在拆分为批次之前是否打乱每个类别的样本。请注意，每个拆分中的样本不会被打乱。
random_stateint, RandomState 实例或 None, 默认为 None: 当 shuffle 为 True 时，random_state 影响索引的排序，从而控制每个类别中每个折叠的随机性。否则，将 random_state 设为 None。传入一个整数可在多次函数调用中获得可重现的输出。请参阅术语表。

另请参阅

RepeatedStratifiedKFold: 重复 Stratified K-Fold n 次。

备注

此实现旨在

生成测试集，使所有测试集包含相同或尽可能接近的类别分布。
对类别标签不变：将 y = ["Happy", "Sad"] 重新标记为 y = [1, 0] 不应改变生成的索引。
当 shuffle=False 时，保留数据集排序中的顺序依赖性：某个测试集中来自类别 k 的所有样本在 y 中是连续的，或者在 y 中被来自非 k 类别的样本分隔。
生成测试集，其中最小和最大的差异至多为一个样本。

0.22 版本中的更改: 先前的实现未遵循最后一个约束。

示例

>>> import numpy as np
>>> from sklearn.model_selection import StratifiedKFold
>>> X = np.array([[1, 2], [3, 4], [1, 2], [3, 4]])
>>> y = np.array([0, 0, 1, 1])
>>> skf = StratifiedKFold(n_splits=2)
>>> skf.get_n_splits(X, y)
2
>>> print(skf)
StratifiedKFold(n_splits=2, random_state=None, shuffle=False)
>>> for i, (train_index, test_index) in enumerate(skf.split(X, y)):
...     print(f"Fold {i}:")
...     print(f"  Train: index={train_index}")
...     print(f"  Test:  index={test_index}")
Fold 0:
  Train: index=[1 3]
  Test:  index=[0 2]
Fold 1:
  Train: index=[0 2]
  Test:  index=[1 3]