I've built a u-net architecture using Keras Functional API but I'm having trouble using the sparse categorical cross entropy loss function. My learning task is multi-class, pixel-wise classification for many 256x256 images. The intended output is a 256x256 mask images with integer values from 0-31 (not every mask will contain each class). I have 32 classes so one-hot encoding gives me an OOM error which is why I don't use categorical cross entropy. The majority of the mask pixels are 0s (which may be part of the problem).
I keep getting loss = nan. I've normalized my input data to have mean = 0, std = 1. If I leave the masks as they are, I get an accuracy around 0.97 and the output masks are all 1s (which is obviously incorrect). If I add 1 to all my masks before performing training, the accuracy is 0. I'm using relu activations with a SoftMax in the last convolutional layer.
It seems the problem likely has to do with the format of my output data, so my main question is, what format should it be in for sparse categorical cross entropy? Should I normalize the mask values to be 0-1? Alternatively, are there any other loss functions or accuracy metrics I can use for training? As far as multi-class classification goes the only function I know of is categorical cross entropy. I can provide additional information about my data, network, etc. if needed.
