Realicemos un sencillo ejemplo de uso de la clase StratifiedKFold a partir del dataset Iris seleccionado, por ejemplo, solo 9 muestras aleatorias extraídas de las clases setosa o virginica:
data = sns.load_dataset("iris")
data = data[data.species.isin(["setosa", "virginica"])]
data["label"] = data.species.astype("category").cat.codes
data = data.sample(9, random_state = 0)
data = data[data.species.isin(["setosa", "virginica"])]
data["label"] = data.species.astype("category").cat.codes
data = data.sample(9, random_state = 0)
data
Extraemos las características predictivas y las etiquetas:
X = data.iloc[:, :4]
y = data.label
y = data.label
Ahora importamos e instanciamos la clase StratifiedKFold especificando tres bloques:
from sklearn.model_selection import StratifiedKFold
scv = StratifiedKFold(n_splits = 3)
Y ya podríamos extraer los índices de los tres grupos de bloques de datos con un bucle:
for train_index, test_index in scv.split(X, y):
print(train_index, test_index)
X_train = X.iloc[train_index]
y_train = y.iloc[train_index]
X_test = X.iloc[test_index]
y_test = y.iloc[test_index]
print(train_index, test_index)
X_train = X.iloc[train_index]
y_train = y.iloc[train_index]
X_test = X.iloc[test_index]
y_test = y.iloc[test_index]
[2 4 5 6 7 8] [0 1 3]
[0 1 3 6 7 8] [2 4 5]
[0 1 2 3 4 5] [6 7 8]
[0 1 3 6 7 8] [2 4 5]
[0 1 2 3 4 5] [6 7 8]
Vemos en el resultado devuelto los índices del conjunto de entrenamiento y, a continuación, los del conjunto de validación, para los tres bloques en los que se ha dividido nuestro dataset.
