В соответствии с предложением @Noki, вы можете использовать параметр cv в CV поиска по сетке.
GridSearchCV(estimator, param_grid, scoring=None, n_jobs=None, iid='deprecated',
refit=True, cv=None, verbose=0,
pre_dispatch='2*n_jobs',error_score=nan,return_train_score=False)
Также в документации четко указано, что если это проблема классификации, она автоматически обеспечит ее стратификацию.
Для входных данных типа integer / None, если оценщик является классификатором, а y является двоичным или мультиклассовым, используется StratifiedKFold. Во всех других случаях используется KFold.
Однако есть кое-что, что я хотел бы добавить: вы можете сделать ваши K-складки динамическими c относительно количества значений вашей Y_target переменная. Вы не можете иметь наименьшее значение вашей частоты в K-кратном значении, равное 1, это приведет к ошибке во время тренировки. Мне довелось столкнуться с этим. Используйте приведенный ниже фрагмент кода, чтобы помочь вам в этом.
Например,
import pandas as pd
Y_target=pd.Series([0,1,1,1,1,0,0,0,6,6,6,6,6,6,6,6,6])
if Y_target.value_counts().iloc[-1]<2:
raise Exception ("No value can have frequency count as 1 in Y-target")
else:
Kfold=Y_target.value_counts().iloc[-1]
Затем вы можете назначить Kfold параметру cv в Grid Search