kerasで自作の損失関数を作ったがLossがnanになる

実現したいこと

LSTMを使って入力データから出力を3つに分類するモデルを作りたいです。モデルの損失関数にClass-Balanced Lossという不均衡データに対応するための計算処理を導入するために、自分で損失関数をカスタムしています。正解ラベルと予測された値はどちらもOne-hotエンコーディングされています。

発生している問題・分からないこと

手動で計算したときはreturnで返ってくるLossの値が存在しているのに対し、実際にモデルを実装して訓練してみるとLossがnanとなってしまい上手く学習できません。計算式が間違っているとも思いにくいのですが、色々試しても解決できませんでした。どなたか力をお貸しください。

エラーメッセージ

error

146/46 [==============================] - 6s 56ms/step - loss: nan - accuracy: 0.1196 - val_loss: nan - val_accuracy: 0.0909 2Epoch 2/10 346/46 [==============================] - 1s 26ms/step - loss: nan - accuracy: 0.1087 - val_loss: nan - val_accuracy: 0.0909 4Epoch 3/10 546/46 [==============================] - 1s 27ms/step - loss: nan - accuracy: 0.1087 - val_loss: nan - val_accuracy: 0.0909 6Epoch 4/10 746/46 [==============================] - 1s 26ms/step - loss: nan - accuracy: 0.1087 - val_loss: nan - val_accuracy: 0.0909

該当のソースコード

def class_balanced_cross_entropy(y_val, y_pred): beta=0.999 y_val = tf.cast(y_val, tf.float32) y_pred = tf.cast(y_pred, tf.float32) y_pred = tf.clip_by_value(y_pred, 1e-6, y_pred) # クラスごとのサンプル数を計算 class_counts = tf.reduce_sum(y_val, axis=0) # クラスごとの重みを計算 class_weights = 1.0 -beta / (1.0 - tf.math.pow(beta, class_counts)) weights = class_weights / tf.reduce_sum(class_weights) * len(class_counts) # 交差エントロピー per_class_loss = -y_val * tf.math.log(y_pred) # クラスごとの損失に重みをかけて合計（Class-Balenced Loss） class_balanced_loss = tf.reduce_sum(per_class_loss * weights,axis=1) # 損失の平均を計算 mean_class_balanced_loss = tf.reduce_mean(class_balanced_loss) return mean_class_balanced_loss def build_model(): model = tf.keras.Sequential([ layers.Input(shape=(50,95)), layers.Masking(mask_value=0), layers.LSTM(units=95,return_sequences=False), layers.Dense(3, activation='softmax') ]) optimizer = Adam(learning_rate=0.001) model.compile(loss=class_balanced_cross_entropy, optimizer=optimizer, metrics=["accuracy"]) return model model = build_model() history = model.fit(X_train, y_train_onehot, epochs=10, batch_size=2, validation_data=(X_test, y_test_onehot))

試したこと・調べたこと

上記の詳細・結果

Class-Balanced Lossを実装している記事をQiitaで見つけましたが、pytorchで実装しており自分なりにtensorflowで書き換えましたが上手くいきません。https://qiita.com/myellow/items/75cd786a051d097efa81

補足

https://openaccess.thecvf.com/content_CVPR_2019/papers/Cui_Class-Balanced_Loss_Based_on_Effective_Number_of_Samples_CVPR_2019_paper.pdf
Class-Balanced Lossを説明している論文です。この論文を元に導入しています。