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]