Generalized Dice Loss: None-gradient when implemented in keras

Question

I'm trying to implement the Generalized Dice Loss as implemented in NiftyNet, in a different convnet written in Keras (TF backend).

def generalised_dice_loss(prediction,
                      ground_truth,
                      weight_map=None,
                      type_weight='Square'):
"""
Function to calculate the Generalised Dice Loss defined in
    Sudre, C. et. al. (2017) Generalised Dice overlap as a deep learning
    loss function for highly unbalanced segmentations. DLMIA 2017

:param prediction: the logits
:param ground_truth: the segmentation ground truth
:param weight_map:
:param type_weight: type of weighting allowed between labels (choice
    between Square (square of inverse of volume),
    Simple (inverse of volume) and Uniform (no weighting))
:return: the loss
"""
prediction = tf.cast(prediction, tf.float32)
if len(ground_truth.shape) == len(prediction.shape):
    ground_truth = ground_truth[..., -1]
one_hot = labels_to_one_hot(ground_truth, tf.shape(prediction)[-1])

if weight_map is not None:
    n_classes = prediction.shape[1].value
    weight_map_nclasses = tf.reshape(
        tf.tile(weight_map, [n_classes]), prediction.get_shape())
    ref_vol = tf.sparse_reduce_sum(
        weight_map_nclasses * one_hot, reduction_axes=[0])

    intersect = tf.sparse_reduce_sum(
        weight_map_nclasses * one_hot * prediction, reduction_axes=[0])
    seg_vol = tf.reduce_sum(
        tf.multiply(weight_map_nclasses, prediction), 0)
else:
    ref_vol = tf.sparse_reduce_sum(one_hot, reduction_axes=[0])
    intersect = tf.sparse_reduce_sum(one_hot * prediction,
                                     reduction_axes=[0])
    seg_vol = tf.reduce_sum(prediction, 0)
if type_weight == 'Square':
    weights = tf.reciprocal(tf.square(ref_vol))
elif type_weight == 'Simple':
    weights = tf.reciprocal(ref_vol)
elif type_weight == 'Uniform':
    weights = tf.ones_like(ref_vol)
else:
    raise ValueError("The variable type_weight \"{}\""
                     "is not defined.".format(type_weight))
new_weights = tf.where(tf.is_inf(weights), tf.zeros_like(weights), weights)
weights = tf.where(tf.is_inf(weights), tf.ones_like(weights) *
                   tf.reduce_max(new_weights), weights)
generalised_dice_numerator = \
    2 * tf.reduce_sum(tf.multiply(weights, intersect))
generalised_dice_denominator = \
    tf.reduce_sum(tf.multiply(weights, seg_vol + ref_vol)) + 1e-6
generalised_dice_score = \
    generalised_dice_numerator / generalised_dice_denominator
return 1 - generalised_dice_score

However, when I try to train the model using this loss function, the following error is raised:

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.

This does not happen with any other loss functions I try to implement. Does anybody know how to address this issue? As far as I can tell, no non-differentiable operators are present in the code.

Answer 1

Similar issue with keras (custom layers and/or models), solution was to build the layers in the def init() before using them in call()

self.lstm_custom_1 = keras.layers.LSTM(128,batch_input_shape=batch_input_shape, return_sequences=False,stateful=True)
self.lstm_custom_1.build(batch_input_shape)

self.dense_custom_1 = keras.layers.Dense(32,  activation = 'relu')
self.dense_custom_1.build(input_shape=(batch_size, 128))