Is there a way to output a metric with several values in keras?

Question

I'm working on a U-net architecture to perform segmentation in 10 classes. I want to calculate the Dice Coefficient for each class after each epoch.

The output of my network is a stack of the segmentation masks for each class with shape (b_size, rows, cols, num_classes). Over this output, I am computing the Dice Coefficient of each class in the following way:

def dice_metric(ground_truth, prediction):
    # initialize list with dice scores for each category
    dice_score_list = list()
    # get list of tensors with shape (rows, cols)
    ground_truth_unstacked = reshape_ground_truth(ground_truth)
    prediction_unstacked = tf.unstack(prediction, axis=-1)
    for (ground_truth_map, prediction_map) in zip(ground_truth_unstacked, prediction_unstacked):
        # calculate dice score for every class
        dice_i = dice_score(ground_truth_map, prediction_map)
        dice_score_list.append(dice_i)
    return tf.reduce_mean(dice_score_list, axis=[0])

Is there any way that I can print the list of dice scores instead of the mean. So in each epoch the output is:

Epoch 107/200
- 13s - loss: 0.8896 - dice_metric: [dice_class_1, ... dice_class_10] - val_loss: 3.3417 - val_dice_metric: [val_dice_class_1, ... val_dice_class_10]

Keras documentation on Custom Metrics only considers single tensor values (i.e., "Custom metrics can be passed at the compilation step. The function would need to take (y_true, y_pred) as arguments and return a single tensor value."

Is there any way/workaround to output a metric with more than one value?

Answer 1

For keras to output all channels, you will need one metric per channel. You can create a wrapper that takes the index and returns only the desired class:

#calculates dice considering an input with a single class
def dice_single(true,pred):
    true = K.batch_flatten(true)
    pred = K.batch_flatten(pred)
    pred = K.round(pred)

    intersection = K.sum(true * pred, axis=-1)
    true = K.sum(true, axis=-1)
    pred = K.sum(pred, axis=-1)

    return ((2*intersection) + K.epsilon()) / (true + pred + K.epsilon())

def dice_for_class(index):
    def dice_inner(true,pred):

        #get only the desired class
        true = true[:,:,:,index]
        pred = pred[:,:,:,index]

        #return dice per class
        return dice_single(true,pred)
    return dice_inner

Then your metrics in the model will be `metrics = [dice_for_class(i) for i in range(10)]

---

Hint: don't iterate unless it's absolutely necessary.

Example of dice for the ten classes without iteration

def dice_metric(ground_truth, prediction):

    #for metrics, it's good to round predictions:
    prediction = K.round(prediction)

    #intersection and totals per class per batch (considers channels last)
    intersection = ground_truth * prediction
    intersection = K.sum(intersection, axis=[1,2])
    ground_truth = K.sum(ground_truth, axis=[1,2])
    prediction = K.sum(prediciton, axis=[1,2])

    dice = ((2 * intersection) + K.epsilon()) / (ground_truth + prediction + K.epsilon())