TensorflowでlossがNanとなる際に、学習・出力状況を表示したい

#問題・解決したいこと
リンク内容でアドバイスをいただきNanの原因を探しているところです。
現在LossがNanになる問題が発生しています。この際に、lossがNanである場合に勾配情報等関連する情報を表示させたいと考えています。
リンク内容を参考にプログラムを書き換えたのですが、Nanになるタイミングでモデルの情報を取り出したいのですが、条件文で止まる様子もなく、またモデルの出力を表示させる部分もうまくいっておりません。
どのように変えたらよいでしょう。

#実装したコード
問題は、出力の表示させたい部分である

print(logits)

と
Nanである際に、何らかの表示をさせたい以下の部分です。

        if train_loss.result() == jud_Nan:
            tf.print(train_loss.result())

以下はコードの全体

from tensorflow.keras.preprocessing.image import ImageDataGenerator
import matplotlib.pyplot as plt

datagen = ImageDataGenerator(
    rescale=1./255,          #rescale で正規化
    validation_split=0.3     #validation_split で検証用データセット分割可能
    
    )

aug_datagen = ImageDataGenerator(
    rescale = 1./255,          
    rotation_range = 20,       
    width_shift_range = 10,     
    height_shift_range = 10   
    )
BATCH_SIZE=10

train_generator = datagen.flow_from_directory(
    '/content/drive/My Drive/my_data_set/train/',
    target_size=(224, 224),
    class_mode='categorical',
    batch_size=BATCH_SIZE,
    subset='training',
)

#拡張データの学習
aug_train_generator = aug_datagen.flow_from_directory(
    '/content/drive/My Drive/my_data_set/train/',
    target_size=(224, 224),
    class_mode='categorical',
    batch_size=BATCH_SIZE,
    subset='training',
)

val_generator = datagen.flow_from_directory(
    '/content/drive/My Drive/my_data_set/train/',
    target_size=(224, 224),
    class_mode='categorical',
    batch_size=BATCH_SIZE,
    subset='validation'
)
import tensorflow as tf
import datetime
import numpy as np

#デバッグのため NaNを作成
jud_Nan = np.nan

# Load the TensorBoard notebook extension
%load_ext tensorboard

os.chdir('/content/drive/My Drive/my_data_set/')
!rm -rf ./logs/ 

log_dir = "logs/fit/" + datetime.datetime.now().strftime("%Y%m%d-%H%M%S")

#log_dir	logファイルを書き出すディレクトリの指定,
#histogram_freq	Tensorboardのhistogram用データを出力する頻度の指定．histogram_freq=1の場合は毎epochデータが出力される
#write_graph	モデルのグラフを出力するか否かの指定
tensorboard_callback = tf.keras.callbacks.TensorBoard(log_dir=log_dir, histogram_freq=1,write_graph=True)
early_stopping_call_back = tf.keras.callbacks.EarlyStopping(monitor='accuracy', patience=10, verbose=0, mode='auto')

#再現性の向上のためシードを固定する
import numpy as np
import random
import os

def set_seed(seed = 200):
    tf.random.set_seed(seed)

    # optional
    # for numpy.random
    np.random.seed(seed)
    # for built-in random
    random.seed(seed)
    # for hash seed
    os.environ["PYTHONHASHSEED"] = str(seed)

#乱数シード値の固定
set_seed(0)

class MyModel(tf.keras.Model):
    def __init__(self):
      
        super(MyModel, self).__init__()
        self.conv1 = tf.keras.layers.Conv2D(8,(1,1),activation='relu',input_shape = (224, 224 ,3))
        self.conv1_2 = tf.keras.layers.Conv2D(32,(5,5),activation='relu')
        self.conv2 = tf.keras.layers.Conv2D(16,(1,1),activation='relu')
        self.conv2_2 = tf.keras.layers.Conv2D(64,(3,3),activation='relu')
        self.conv3 = tf.keras.layers.Conv2D(32,(1,1),activation='relu')
        self.conv3_2 = tf.keras.layers.Conv2D(128,(3,3),activation='relu')
        self.pooling1 = tf.keras.layers.MaxPooling2D((2,2), strides=None, padding='valid')
        self.pooling2 = tf.keras.layers.MaxPooling2D((2,2), strides=None, padding='valid')
        self.flatten = tf.keras.layers.Flatten()
        self.fc1 = tf.keras.layers.Dense(256, activation='relu')
        self.fc2 = tf.keras.layers.Dense(9, activation='relu')
        self.dropout = tf.keras.layers.Dropout(0.2)

    def call(self, x, training=False):
        x = self.conv1(x)
        x = self.conv1_2(x)
        x = self.pooling1(x)
        x =  self.conv2(x)
        x =  self.conv2_2(x)
        x = self.pooling2(x)
        x =  self.conv3(x)
        x =  self.conv3_2(x)
        x = self.pooling2(x)       
        x = self.flatten(x)
        x = self.fc1(x)
        x = self.dropout(x, training=training)
        x = self.fc2(x)
        return x

# モデルインスタンスの生成
model = MyModel()

# tf.functionでグラフへコンパイルする
@tf.function
def train_step(images, labels):
    with tf.GradientTape() as tape: 
        logits = model(images) 

        #毎ステップ表示されるようにしたい 内容を表示されるようにしたい
        print(logits)

        loss_value = loss_object(y_true = labels, y_pred = logits)

    # 損失に対する勾配を求める 
    grads = tape.gradient(loss_value, model.trainable_variables)
    # 指定した最適化関数を用いてモデルの更新(zipを忘れないように注意)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))
    # 訓練の損失を記録 train_lossグローバル変数
    train_loss(loss_value)
    # 訓練の正答率を記録 train_accuracy グローバル変数
    train_accuracy(y_true=labels, y_pred=logits)

@tf.function
def test_step(images, labels):
    # モデルの前進処理 
    predictions = model(images)
    # 損失の計算
    t_loss = loss_object(labels, predictions)
    # テストの損失を記録 test_lossグローバル変数
    test_loss(t_loss)
    # テストの正答率を記録 test_lossグローバル変数
    test_accuracy(labels, predictions)

loss_object = tf.keras.losses.CategoricalCrossentropy()

optimizer = tf.keras.optimizers.Adam()

# 訓練の損失を記録するインスタンスの生成
train_loss = tf.keras.metrics.Mean(name='train_loss')
# 訓練の正答率を記録するインスタンスの生成
train_accuracy = tf.keras.metrics.CategoricalAccuracy(name='train_accuracy')
# テストの損失を記録するインスタンスの生成
test_loss = tf.keras.metrics.Mean(name='test_loss')
# テストの正答率を記録するインスタンスの生成
test_accuracy = tf.keras.metrics.CategoricalAccuracy(name='test_accuracy')

# model.fitに代わる関数
# @tf.function デバッグが難しい
def train(model, x_train, t_train, epochs, batch_size):
    # Epoch回train_stepで訓練する 
    for epoch in range(epochs):
        train_step(x_train,t_train)
        # 訓練履歴の出力
        template = 'Epoch {}, Loss: {}, Accuracy: {}'
        tf.print(template.format(epoch+1,
                              train_loss.result(),
                              train_accuracy.result()*100))
        
        #Nanの際に条件を満たしているのに表示されない
        if train_loss.result() == jud_Nan:
            tf.print(train_loss.result())


    # 訓練履歴のリセット
    train_loss.reset_states()
    train_accuracy.reset_states()

train(model,training_images , training_labels, epochs=25, batch_size=1)

#試したこと
上記のコード内でprint文での情報の表示を試みました。

(追記 2020/10/5)
リンク内容を参考にデコレートされた関数をtf.Printで、また、train_loss.result()に関しては、リンク内容を参考に.numpy()をつけることで解決いたしました。しかし、条件文に関してはいまだ解決していません。
引き続き回答を募集いたします。

#環境
googlecolab
tensorflow 2.3.0

行動規範の内容に同意します

回答1件

自己解決

各解決法について以下に示す
デコレートされた関数内部ではtf.Print(logits)とすることで解決した。
train_loss.result()は.numpy()を使いtrain_loss.result().numpy()とすることで解決した。
条件部分は、条件をtf.math.is_nanとすることで解決した。

投稿2020/10/05 11:19

radbag

総合スコア0