sklearnのStratifiedKFoldの実装の注意点について

#コード
実装例は次のようになる。

from sklearn import model_selection
import numpy as np
#説明変数X
X = np.array([['A','B'],['C','D'],['E','F'],['G','H']])

#目的変数y
y = np.array([0,0,1,1])

kf = model_selection.StratifiedKFold(n_splits = 2)

for fold,(train_index, test_index) in enumerate(kf.split(X=X,y=y)):
 print("kfold:",fold)
 print(f"train:{X[train_index]}")
 print(f"test:{X[test_index]}")

出力は次のようになる。

kfold: 0
train:[['C' 'D']
 ['G' 'H']]
test:[['A' 'B']
 ['E' 'F']]
kfold: 1
train:[['A' 'B']
 ['E' 'F']]
test:[['C' 'D']
 ['G' 'H']]

#注意点
kf.splitの戻り値は渡したデータのインデックスであること。上のコードならk=0のときtrainのインデックスは1,3で、testのインデックスは0,2である。

このインデックスを用いて目的変数を参照するとy[1]=0,y[3]=1かつy[0]=0,y[2]=1であるのできちんと割合一定で分割ができていることが確認できる。

fold数は出力しないのでenumerate関数を用いてこちらで番号を付けている。

#参考にさせたいただいたサイト
sklearnの交差検証の種類とその動作
Python, enumerateの使い方: リストの要素とインデックスを取得