Custom mask loss giving error

[Paulito-7]

hello everyone,

I am currently trying to change the mask loss for weighted bce dice loss available here : https://www.kaggle.com/lyakaap/weighing-boundary-pixels-loss-script-by-keras2?scriptVersionId=1460731

To make it compatible with this implementation, I took the same gathering of prediction and ground truth masks as follows :

def weighted_bce_dice_loss(target_masks, target_class_ids, pred_masks):
    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    y_true = tf.gather(target_masks, positive_ix)
    y_pred = tf.gather_nd(pred_masks, indices)

    y_true = K.cast(y_true, 'float32')
    y_pred = K.cast(y_pred, 'float32')
    # if we want to get same size of output, kernel size must be odd number
    averaged_mask = K.pool2d(
            y_true, pool_size=(11, 11), strides=(1, 1), padding='same', pool_mode='avg')
    border = K.cast(K.greater(averaged_mask, 0.005), 'float32') * K.cast(K.less(averaged_mask, 0.995), 'float32')
    weight = K.ones_like(averaged_mask)
    w0 = K.sum(weight)
    weight += border * 2
    w1 = K.sum(weight)
    weight *= (w0 / w1)
    loss = weighted_bce_loss(y_true, y_pred, weight) + \
    weighted_dice_loss(y_true, y_pred, weight)
    return loss 

weighted_bce_loss and weighted_dice_loss being exactly the same as on the link above.

And I changed in model.py --> class MaskRCNN -->build :

mask_loss = KL.Lambda(lambda x: mrcnn_mask_loss_graph(*x), name="mrcnn_mask_loss")([target_mask, target_class_ids, mrcnn_mask])`

To

mask_loss = KL.Lambda(lambda x: weighted_bce_dice_loss(*x), name="mrcnn_mask_loss")([target_mask, target_class_ids, mrcnn_mask])

Yet I get the following error :

ValueError: Shape must be rank 4 but is rank 3 for 'mrcnn_mask_loss/AvgPool' (op: 'AvgPool') with input shapes: [?,?,?].

Does anyone know what to do to figure it out? I feel a bit lost..

Thank you!

[Paulito-7]

I managed to adapt it for the mask R-CNN, for those interested, here is the code of the loss :

def weighted_bce_loss(y_true, y_pred, weight):
    # avoiding overflow
    epsilon = 1e-7
    y_pred = K.clip(y_pred, epsilon, 1. - epsilon)
    logit_y_pred = K.log(y_pred / (1. - y_pred))

    loss = (1. - y_true) * logit_y_pred + (1. + (weight - 1.) * y_true) * \
    (K.log(1. + K.exp(-K.abs(logit_y_pred))) + K.maximum(-logit_y_pred, 0.))
    return K.sum(loss) / K.sum(weight)

def weighted_dice_loss(y_true, y_pred, weight):
    smooth = 1.
    w, m1, m2 = weight * weight, y_true, y_pred
    intersection = (m1 * m2)
    score = (2. * K.sum(w * intersection) + smooth) / (K.sum(w * m1) + K.sum(w * m2) + smooth)
    loss = 1. - K.sum(score)
    return loss

def weighted_bce_dice_loss(target_masks, target_class_ids, pred_masks):

    # Reshape for simplicity. Merge first two dimensions into one.
    target_class_ids = K.reshape(target_class_ids, (-1,))
    mask_shape = tf.shape(target_masks)
    target_masks = K.reshape(target_masks, (-1, mask_shape[2], mask_shape[3]))
    pred_shape = tf.shape(pred_masks)
    pred_masks = K.reshape(pred_masks,
                           (-1, pred_shape[2], pred_shape[3], pred_shape[4]))
    # Permute predicted masks to [N, num_classes, height, width]
    pred_masks = tf.transpose(pred_masks, [0, 3, 1, 2])

    # Only positive ROIs contribute to the loss. And only
    # the class specific mask of each ROI.
    positive_ix = tf.where(target_class_ids > 0)[:, 0]
    positive_class_ids = tf.cast(
        tf.gather(target_class_ids, positive_ix), tf.int64)
    indices = tf.stack([positive_ix, positive_class_ids], axis=1)

    # Gather the masks (predicted and true) that contribute to loss
    target_masks = tf.gather(target_masks, positive_ix)
    pred_masks = tf.gather_nd(pred_masks, indices)

    target_masks = K.expand_dims(target_masks,0)

    pred_masks = K.expand_dims(pred_masks,0)

    y_true = K.cast(target_masks, 'float32')
    y_pred = K.cast(pred_masks, 'float32')

    averaged_mask = K.pool2d(
            y_true, pool_size=(11, 11), strides=(1, 1), padding='same', pool_mode='avg')
    border = K.cast(K.greater(averaged_mask, 0.005), 'float32') * K.cast(K.less(averaged_mask, 0.995), 'float32')
    weight = K.ones_like(averaged_mask)
    w0 = K.sum(weight)
    weight += border * 2
    w1 = K.sum(weight)
    weight *= (w0 / w1)
    loss = weighted_bce_loss(y_true, y_pred, weight) + \
    weighted_dice_loss(y_true, y_pred, weight)
    return loss 

[vaibhav0195]

hi @Paulito-7 i was trying to integrate your loss for my custom dataset. but i am getting nan values for the bce loss . the output of my iteration is

---

3/12054 [..............................] - ETA: 13:16:07 - loss: nan - rpn_class_loss: 0.6509 - rpn_bbox_loss: 4.7030 - mrcnn_class_loss: 0.1926 - mrcnn_bbox_loss: 0.3745 - mrcnn_mask_loss: nan - bceMaskLoss: nan - dceMaskLoss: 0.2165 - summOfLoss: 682.5248 - weightLoss: 784.0000

here bceMaskLoss is the bce loss calculated to get the maskloss ,summOfLoss is the sum of the loss inside the bce-loss and weightLoss is the sum of the weight loss.

---

its strange that both the sumOfLoss and weightLoss are not nan but the final BCE loss is calculated as nan.

any help would be appericiated. thanks