kerasの独自レイヤを作成したが、lossがnanになってしまう。

Question

いつもお世話になっております。
現在、kaggleのコンペでどうしてもLogSoftmaxレイヤをkerasで実装したくて、以下のコードを書き、InspectionResNetでファインチューニングさせているところです。しかし、lossの値がnanと表示されていて、学習もうまくいっていないようです。
調べてもよくわからないので、ご教授いただけると嬉しいです。

from tensorflow.keras import layers
from tensorflow.keras import activations
class LogSoftmax(layers.Layer):
  def __init__(self, axis=-1, **kwargs):
      super(LogSoftmax, self).__init__(**kwargs)
      self.supports_masking = True
      self.axis = axis
  def call(self,inputs):
    softmax = activations.softmax(inputs, axis=self.axis)
    return tf.math.log(softmax)
  def get_config(self):
      config = {'axis': self.axis}
      base_config = super(Softmax, self).get_config()
      return dict(list(base_config.items()) + list(config.items()))
  def compute_output_shape(self, input_shape):
      return input_shape

以下の部分でモデルをコンパイルして訓練しています。

model = tf.keras.Sequential([InceptionResNetV2(input_shape=(512, 512, 3),
                                          include_top=False,weights="imagenet"),
                              L.GlobalMaxPooling2D(),
                              L.Dense(4)])
model.add(LogSoftmax())
model.compile(optimizer='adam',
              loss = 'categorical_crossentropy',
              metrics=['categorical_accuracy'])
model.summary()
x_train,x_test=train_data.iloc[train_index],train_data.iloc[test_index]
  train_datagenerator = datagen.flow_from_dataframe(
    directory = IMAGE_PATH,
    x_col = "～～～",
    y_col = "～～～",
    dataframe=x_train.astype("str"),
    )
  valid_datagenerator = datagen.flow_from_dataframe(
    directory = IMAGE_PATH,
    x_col = "～～～",
    y_col = "～～～",
    dataframe=x_test.astype("str"),
    )
history = model.fit(train_datagenerator, steps_per_epoch=len(train_datagenerator),epochs=2,validation_steps=len(valid_datagenerator),validation_data=valid_datagenerator,callbacks=[es_cb])

定義したクラスの概要としては、logSoftmaxクラス内では、githubに上がっていた、Softmaxレイヤのcall関数内を書き換えました。
具体的には、softmaxで計算したテンソルのlogを返すように仕様変更しました。
環境はgoogle colabです。
長文になって申し訳ないのですが、ご教授いただけると嬉しいです。

Comments

[meg_]

コンペは終了しているものですか？

[umepon0626]

いえ、現在進行中のものです。https://www.kaggle.com/c/plant-pathology-2020-fgvc7

disscussionの中で、logsoftmaxを使うといいよ！ってのがあったのですが、notebookを探しても実装を見つけることができませんでした。Pytorchならあるっぽいのですが、現在、モデルの大半をkerasで実装してしまっているので、kerasのレイヤを自作しようと思った次第です。

[meg_]

この質問は下記規約に違反していませんか？

No private sharing outside teams
Privately sharing code or data outside of teams is not permitted. It's okay to share code if made available to all participants on the forums.

[umepon0626]

データエンジニアリングや画像データに関することは規約に違反すると思いましたが、今回の質問の場合、コンペというよりはkerasに関する質問という認識だったので質問させていただいたのですが、規約違反となるのでしょうか？