TF.Keras SparseCategoricalCrossEntropy 在 GPU 上返回 nan

Question

Tried to train UNet on GPU to create binary classified image. Got nan loss on each epoch. Testing of loss function always produces nan-return.

Test case:

import tensorflow as tf
import tensorflow.keras.losses as ls

true = [0.0, 1.0]
pred = [[0.1,0.9],[0.0,1.0]]

tt = tf.convert_to_tensor(true)
tp = tf.convert_to_tensor(pred)

l = ls.SparseCategoricalCrossentropy(from_logits = True)
ret = l(tt,tp)

print(ret) #tf.Tensor(nan, shape=(), dtype=float32)

If i would force my tf to work with CPU ( Can Keras with Tensorflow backend be forced to use CPU or GPU at will? ), all works fine. And yes, my UNet fits and predicts correctly on CPU.

I checked several posts on keras GitHub, but the all point to problems with compiled ANN, such as using inappropriate optimizers for categorical crossentropy.

Any workaround? Am i missing something?

Answer 1

The test code you have provided is working fine on google colab.

tf.__version__

2.3

tf.config.list_physical_devices('GPU')  

Output:

[PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]  

Your code:

import tensorflow as tf
import tensorflow.keras.losses as ls

true = [0.0, 1.0]
pred = [[0.1,0.9],[0.0,1.0]]

tt = tf.convert_to_tensor(true)
tp = tf.convert_to_tensor(pred)

l = ls.SparseCategoricalCrossentropy(from_logits = True)
ret = l(tt,tp)

print(ret)  

Result:

tf.Tensor(0.8132616, shape=(), dtype=float32)

Answer 2

I had the same issue. My loss was real number if I trained on CPU. I tried to upgrade tf version but it wasn't the fix. I finally fixed my issue by reduce the y dimension. My model output was 2D array. when I reduce to 1D, I can get real loss on GPU.