In order to use a binary segmentation loss function in a multi label problem, I would like to permute the batch axis with the channel axis in the loss computation in order to compute the loss by channel rather than batches to fight a class inbalance problem. (for example, if a class A is very predominant on the others, I could obtain a good loss result by attributing class A to each pixels / voxels.)

Here is an example using the dice loss on 3d images:

def dice_coef(y_true, y_pred, smooth=1e-10):

    # permute channels with batches
    # eg:  [num_batches, height, width, depth, num_channels] --becomes--> [num_channels, height, width, depth, num_batches]
    # in order to perform the computation per channel rather than batches
    y_true = K.permute_dimensions(y_true, (4,1,2,3,0))
    y_pred = K.permute_dimensions(y_pred, (4,1,2,3,0))

    y_true_f = K.batch_flatten(y_true)
    y_pred_f = K.batch_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 1 - dice_coef(y_true, y_pred)

is my raisonning correct ?


Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.