Custom Loss Function does not equal Loss Function used to train model

Viewed 101

this is my first post so I will try to detail all the relevant info. If anything is missing please let me know!

I am currently trying to create cnn (based off unet) for image segmentation on grayscale images.

I have created a custom function to calculate dice loss and Binary Cross Entropy loss, see below.

def dice_BCE_coef_loss(y_true, y_pred):
    smooth = 1
    bce_weight = 0.5

    #y_true_f = tensorflow.math.reduce_sum(y_true)
    #y_pred_f = tensorflow.math.reduce_sum(y_pred)
    y_true_f = tensorflow.reshape(y_true, [-1])
    y_pred_f = tensorflow.reshape(y_pred, [-1])

    intersection  = tensorflow.math.reduce_sum(y_true_f * y_pred_f)
    union = tensorflow.math.reduce_sum(y_true_f + y_pred_f)
    dice_coef = (2*intersection + smooth) / (union + smooth)
    dice_loss = 1 - dice_coef

    BCE =  tensorflow.keras.losses.BinaryCrossentropy(from_logits=True)(y_true, y_pred)
    dice_BCE = tensorflow.math.reduce_mean(BCE * bce_weight + dice_loss * (1 - bce_weight))
    return dice_BCE

I then add this to my model as the loss.

model.compile(optimizer=tensorflow.keras.optimizers.Adam(lr=1e-3), 
              loss=dice_BCE_coef_loss, 
              metrics=['accuracy']
              )

The issue comes when I calculate the dice_BCE manually the loss value is different to the output loss during training. To confirm whether this was a correct value across the whole dataset (my manual check was a single image) I reduced my dataset to a single image and mask yet they still didn't match.

Image showing the discrepancy of loss vs my expected dice_BCE loss (I hope a picture is allowed in this case) 1

This loss always remains around 0.48 after several epochs, however never really improves from there and sometimes you can see the output mask is really close (and a good expected dice_BCE to match) yet it ends up diverging because the loss it seems to train on can be improved in other ways (but increase the expected dice_loss).

The dice loss (through the loss value of the epoch) is also far lower than when calculated through the function. Around 0.001 even when the accuracy of the prediction is ~50% and it visible looks incorrect.

Can anybody explain how this loss is calculated and why it does not match what I expect it to be?

I have read through similar posts on here but can't find anything useful.

Any suggestions of what to look into next or resources to investigate further if this is obvious please do let me know! Thank you in advance

0 Answers
Related