How can I fix my dice loss calculation during the training process of a convolutional neural network implemented in Keras?

Question

I am implementing my own code using keras to do semantic segmentation. My testing images has shape (10, 512, 512, 5), where 10 is the number of images, 512 is their size and 5 is the number of classes I want to segment. As last activation function I use softmax and as loss I want to extract the dice loss (https://arxiv.org/abs/1606.04797) in order to improve the segmentation results. My code is:

eps = 1e-3

def dice(y_true, y_pred):
    y_pred = K.one_hot(K.argmax(y_pred,axis=-1), Nclasses) 
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    num = 2*K.sum(y_true_f*y_pred_f)
    den = K.sum(K.square(y_true_f))+K.sum(K.square(y_pred_f))+eps
    return num/den

def dice_loss(y_true, y_pred):
    return 1-dice(y_true, y_pred)

I use K.one_hot(K.argmax(...)) because in this way my y_pred is binary and not made by probabilities (right?). Anyway, when the training process starts, I receive this error:

"ValueError: An operation has None for gradient. Please make sure that all of your ops have a gradient defined (i.e. are differentiable). Common ops without gradient: K.argmax, K.round, K.eval."

Answer 1

Try using this code snippet for your dice coefficient. Important observation : If you have your masks one-hot-encoded, this code should also work for multi-class segmentation.

smooth = 1.

def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2. * intersection + smooth) / (K.sum(y_true_f) + K.sum(y_pred_f) + smooth)


def dice_coef_loss(y_true, y_pred):
    return -dice_coef(y_true, y_pred)