How to do point-wise categorical crossentropy loss in Keras?

Question

I have a network that produces a 4D output tensor where the value at each position in spatial dimensions (~pixel) is to be interpreted as the class probabilities for that po

Answer 1

It seems that now you can simply do softmax activation on the last Conv2D layer and then specify categorical_crossentropy loss and train on the image without any reshaping tricks or any new loss function. I've tried overfitting with a dummy dataset and it works well. Try it ~ !

inp = keras.Input(...)
# define your model here
out = keras.layers.Conv2D(classes, (1, 1), activation='softmax') (...)
model = keras.Model(inputs=[inp], outputs=[out], name='unet')
model.compile(loss='categorical_crossentropy',
                      optimizer='adam',
                      metrics=['accuracy'])
model.fit(tensor4d, tensor4d)

You can also compile using sparse_categorical_crossentropy and then train with output of shape (samples, height, width) where each pixel in the output corresponds to a class label: model.fit(tensor4d, tensor3d)

The idea is that softmax and categorical_crossentropy will be applied to the last axis (you can check keras.backend.softmax and keras.backend.categorical_crossentropy doc).

PS. I use keras from tensorflow.keras (tensorflow 2)

Update: I have trained on my real dataset and it is working as well.