Improve Performance of categorical crossentropy on four dimensional prediction and target

Question

I have a Neural Network implemented in tensorflow that takes a batch of images and predicts a tensor with shape (batch_size, 9, 9, 10)

For each image in the batch the target (and the prediction) looks like this. We predict an image with width 9, height 9 and depth 10. The depth of every pixel corresponds to a class so we have 10 classes.

So e.g. the prediction of pixel(0,0) with softmax looks like this [0, 0.8 ,0 ,0 ,0 ,0 ,0.2 ,0 ,0,0] and the target of the pixel(0,0) looks like this [0, 1,0,0,0,0,0,0,0]

I want to use the classical cross entropy between these two tensors.

currently my loss is implemented like this:

def calculate_loss(prediction, target):
    '''
    prediction and target have shape (batch_size, 9, 9, 10)
    '''
    loss = 0
    for batch in range(0, BATCH_SIZE):
        for width in range(0,9):
            for heigth in range(0,9):
                loss += tf.losses.categorical_crossentropy(target[batch][width][height] , prediction[batch][width][height])
return loss

But with a high batch size (1024) this calculation takes 20 seconds. And this is just way too slow. Are there any already implemented functions in tensorflow/keras that can improve the performance of this loss? Or do you know if I can change the code to improve performance?

relot relot · Accepted Answer · 2019-12-23T14:43:23

tf.nn.softmax_cross_entropy_with_logits(target, prediction)

Is the answer. I can then reduce sum the values of the output to get the same thing. Don't forget to remove the sofmax layer in your model.

Improve Performance of categorical crossentropy on four dimensional prediction and target

1 Answers