ニューラルネットワーク(keras)の分割交差検証(KFold)とoptunaの設定について

ニューラルネットワーク(keras)の分割交差検証(KFold)とoptunaの設定について質問です。

ニューラルネットワークにて3分類(0,1,2)を行うプログラムを構築しております。
最適なDenseやunitsを見つけるのにoptunaを利用しています。

分割交差検証を行いたいと考えているのですがどの様に書くと良いでしょうか?

StratifiedKFoldを用いて交差検証行いたいと考えています。

該当箇所のコード


input_dim = X_train.shape[1] #input_shape input_dim

def create_model(n_layer, activation, mid_units, dropout_rate):
    model = keras.Sequential()

    for i in range(n_layer):
        if i == 0:
            model.add(keras.layers.Dense(units = mid_units, activation = "tanh", input_dim = input_dim))
        else:
            model.add(keras.layers.Dense(units = mid_units, activation = "tanh"))
        model.add(keras.layers.Dropout(dropout_rate))
        model.add(keras.layers.BatchNormalization())

    model.add(keras.layers.Dense(3, activation='softmax'))
    return model

def objective(trial):
    #セッションのクリア
    keras.backend.clear_session()

    n_layer      = trial.suggest_int('n_layer', 1, 3) # 追加する層を1-3から選ぶ
    mid_units    = int(trial.suggest_discrete_uniform('mid_units', 32, 300, 1)) # ユニット数
    dropout_rate = trial.suggest_uniform('dropout_rate', 0, 1) # ドロップアウト率
    activation   = trial.suggest_categorical('activation', ['tanh']) # 活性化関数
    optimizer    = trial.suggest_categorical('optimizer', ['adam', 'Adamax', 'Nadam']) # 最適化アルゴリズム

    model = create_model(n_layer, activation, mid_units, dropout_rate)
    model.compile(optimizer=optimizer,
          loss="sparse_categorical_crossentropy",
          metrics=["accuracy"])


    # 雰囲気的にはここらへんいStratifiedKFoldを入れる感じでしょうか?
    # kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
    # scores = cross_validate(model, X_train, y_train, cv=kf)

    history = model.fit(X_train, y_train, epochs = 5, validation_split = 0.1, verbose = 0, batch_size=128)
    return -np.amax(history.history["val_accuracy"])

study = optuna.create_study()
study.optimize(objective, n_trials = 100)

best_params = study.best_params
print("best_params: ", best_params)

なんとなくですがmodel.fitさせる前にStratifiedKFoldの入れて値を返すのかなーと思うのですがいかがでしょうか?
ランダムフォレストやxgboostなどの機械学習ではこ上記の様な方法で行っていました。
実際やっていみるとSequentialにscoresがないのでエラーになります。

または以下ブログを参考に最初の段階で分割してforでその分割分回すのが一般的になりますか?
https://manareki.com/k_fold

この場合は

kf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)

def objective(trial):
    for train_index, val_index in kf.split(X_train,Y_train):

        train_data=X_train[train_index]
        train_label=Y_train[train_index]
        val_data=X_train[val_index]
        val_label=Y_train[val_index]
    
        keras.backend.clear_session()

        n_layer      = trial.suggest_int('n_layer', 1, 3) # 追加する層を1-3から選ぶ
        mid_units    = int(trial.suggest_discrete_uniform('mid_units', 32, 300, 1)) # ユニット数
        dropout_rate = trial.suggest_uniform('dropout_rate', 0, 1) # ドロップアウト率
        activation   = trial.suggest_categorical('activation', ['tanh']) # 活性化関数
        optimizer    = trial.suggest_categorical('optimizer', ['adam', 'Adamax', 'Nadam']) # 最適化アルゴリズム

        model = create_model(n_layer, activation, mid_units, dropout_rate)
        model.compile(optimizer=optimizer,
              loss="sparse_categorical_crossentropy",
              metrics=["accuracy"])


        history = model.fit(train_data, train_label, epochs = 5, validation_split = 0.1, verbose = 0, batch_size=128)
        return -np.amax(history.history["val_accuracy"])

の様な書き方になりますでしょうか?

よろしくお願いします。

tensorflow 2.4.0
optuna 2.1.0
Keras 2.3.1
windows 10
python 3.7