Implement the following loss function without interrupting the gradient chain registered by the gradient tape
Your intuition about counting the 2 is pretty good. For this type of problem, you can use the dice loss function wich use your idea but sligthly differently.



the values of predicted and ground truth components (each pixels) are either 0 or 1, representing whether the pixel bayond to the label (value of 1) or not (value of 0). Therefore, the denominator is the sum of total labeled pixels of both prediction and ground truth, and the numerator is the sum of correctly predicted boundary pixels because the sum increments only when predicted and ground truth match (both of value 1).

as cited here:

For example, if two sets A and B overlap perfectly, DSC gets its maximum value to 1. Otherwise, DSC starts to decrease, getting to its minimum value to 0 if the two sets don ‘t overlap at all. Therefore, the range of DSC is between 0 and 1, the larger the better. Thus we can use 1-DSC as Dice loss to maximize the overlap between two sets.

The dice loss is commonly used in segmentation task because it is resilien to class inbalance (eg too much background)

you can implement it in tensorflow keras like this;

def dice_coef(y_true, y_pred):
    y_true_f = K.flatten(y_true)
    y_pred_f = K.flatten(y_pred)
    intersection = K.sum(y_true_f * y_pred_f)
    return (2.0 * intersection + 1.0) / (K.sum(y_true_f) + K.sum(y_pred_f) + 1.0)

def dice_coef_loss(y_true, y_pred):
    return 1-dice_coef(y_true, y_pred)


(in the implementation, we use a smoothing termes (here it's 1, the lower the value is the better) The smoothing therme is here to avoid dividing by 0 if the input as no label in it)