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 ?