softmax and categorical cross entropy for 4d tensors

[Vladimir Iglovikov]

I am trying to work on the image segmentation problem.

Each pixel in the target image may belong to one of the N classes. Thus, similar to usual classification problem I convert each pixel of the target image to an array of length N, where one of the values in this array is equal to 1 and others are zeros.

=> Target becomes 4d tensor (batch_size, num_classes, img_rows, img_cols)

[1] How may I define softmax activation function that will work for this case? 
[2] How may I define cross_entropy for this case?

Currently, in keras softmax is defined in for 2D and 3D tensors as  https://github.com/fchollet/keras/blob/master/keras/activations.py

Can anyone help me to extend it to 4D? Or may be I am missing something and this loss function for image segmentation does not even worth considering?


Thank you. 

Vladimir

[mike.mi...@gmail.com]

This might help https://github.com/heuritech/convnets-keras/blob/master/convnetskeras/customlayers.py

Also you can reshape, softmax, then reshape back if you like.