Python sklearn StratifiedKFold用法及代码示例

本文简要介绍python语言中 sklearn.model_selection.StratifiedKFold 的用法。

用法:

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

分层 K-Folds cross-validator。

提供训练/测试索引以拆分训练/测试集中的数据。

此交叉验证对象是返回分层折叠的 KFold 的变体。通过保留每个类别的样本百分比来进行折叠。

在用户指南中阅读更多信息。

参数：

n_splits：整数，默认=5

折叠次数。必须至少为 2。

shuffle：布尔，默认=假

是否在分成批次之前对每个类的样本进行洗牌。请注意，每个拆分中的样本不会被打乱。

random_state：int、RandomState 实例或无，默认=无

当 shuffle 为 True 时，random_state 会影响索引的顺序，从而控制每个类的每次折叠的随机性。否则，将 random_state 保留为 None 。传递 int 以在多个函数调用之间实现可重现的输出。请参阅术语表。

注意：

该实施旨在：

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

例子：

>>> 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 train_index, test_index in skf.split(X, y):
...     print("TRAIN:", train_index, "TEST:", test_index)
...     X_train, X_test = X[train_index], X[test_index]
...     y_train, y_test = y[train_index], y[test_index]
TRAIN: [1 3] TEST: [0 2]
TRAIN: [0 2] TEST: [1 3]