Yラベルに偏りのある多クラス分類に於いてoptunaで最適化したモデルをload

解きたい課題

多クラス分類課題に対し、kerasを用いてDNNで学習し、optunaで最適化したモデルを用いてテストデータに対する予測を行おうとしています。
その中で教師データのラベリングに偏りがあり("No1"と"No2"と"No3"と"No4"があるのですが、9割が"No1")、metricsとしてf1_scoreを使おうとしています。

コードの構成

コードの仕組みとしては
optunaのobjective()の中で、すべてのtrialにおいて毎度モデルをtrial.numberで名前を付けてsaveし、
最終的にベストだとされたtrial.numberでそれをload_model()しようとしています。

f1_scoreはmetrics=[f1]などとしても呼び出せなかったので、
別途、def f1(y_true, y_pred)と関数を作って呼び出し、
optunaで最大化する対象をhistory.history["val_f1"][-1]としています。
accuracyなどのメトリクスではテストデータに対する予測が全て"No1"となってしまうため、クラスごとの予測精度を重視するメトリクスを使いたいです
(まずはある程度精度が出ればf1_scoreでなくてもいいのですが、recallなどでは全て"No1"を出力してしまったため、f1scoreを実装しようとしています)

困っていること、エラーログ

optunaでの最適化が終わった後、テストデータに対する予測を行うためload_model()としたところで
モデルのコンパイル時に設定したf1のメトリクスが読みだせないとのことでエラーが出ているのだと思います。
まずはエラーを回避したいのですが、そもそもコードの構成として「こういう風にしたらいいよ」というアドバイスがありましたら是非戴きたいです。
そもそもf1_scoreって自作しないといけないんですかね？

error
1ValueError: Unable to restore custom object of type _tf_keras_metric. Please make sure that any custom layers are included in the 'custom_objects' arg when calling 'load_model()' and make sure that all layers implement 'get_config' and 'from_config'.

バージョン情報

keras 2.7.0
tensorflow 2.7.0

具体的なコードの一部

python
1import tensorflow as tf
2import keras.backend as K
3import optuna
4
5def f1(y_true, y_pred):
6    y_pred = K.round(y_pred)
7    tp = K.sum(K.cast(y_true*y_pred, 'float'), axis=0)
8    # tn = K.sum(K.cast((1-y_true)*(1-y_pred), 'float'), axis=0)
9    fp = K.sum(K.cast((1-y_true)*y_pred, 'float'), axis=0)
10    fn = K.sum(K.cast(y_true*(1-y_pred), 'float'), axis=0)
11
12    p = tp / (tp + fp + K.epsilon())
13    r = tp / (tp + fn + K.epsilon())
14
15    f1 = 2*p*r / (p+r+K.epsilon())
16    f1 = tf.where(tf.math.is_nan(f1), tf.zeros_like(f1), f1)
17    return K.mean(f1)
18
19def savemodel(model, num):
20    os.makedirs("models/", exist_ok=True)
21    model.save("models/" + f"model{num}.h5")
22
23def loadmodel(num):
24    model = load_model("models/" + f"model{num}.h5")
25    return model
26
27def hyperparasearch(X, Y):
28    def objective(trial):
29        #下記ネットワークとハイパラは実際にはもっと複雑ですが単純化しています
30        nodes_num = trial.suggest_int("node", 20, 200)
31        inputs1 = Input(shape=(X.shape[1],))
32        midlayer = Dense(nodes_num, activation="relu")(inputs1)
33        outputs1 = Dense(4, activation="softmax")(midlayer)
34        model = Model(inputs1, outputs1)
35
36        model.compile("adam", loss="categorical_crossentropy", metrics=[f1])
37        history1 = model.fit(X, Y, epochs=100, validation_split=0.2)
38        
39        #モデルを都度保存
40        savemodel(model, trial.number)
41        return history.history['val_f1'][-1]
42    return objective
43
44def main():
45    sampler = optuna.samplers.TPESampler(seed=42)
46    X, Y = get_train_dataset() #単に学習にかけるXとYを取ってきます。YはOne-hot-encodingされています
47    study = optuna.create_study(sampler=sampler, direction="maximize")
48    objective = hyperparasearch(X, Y)
49    study.optimize(objective, n_trials=100)
50
51    best_model = loadmodel(study.best_trial.number) #ここでエラーが出ます