Keras-Tensorflow ValueError: during model.fit whehn using Jaccard or IOU

Question

I am starting my journey in deep learning and relatively new to this topic. Hoping someone could help me out as I am stuck here for a long time.

I am trying to train deeplabv3+ architecture-based model for semantic segmentation when I faced the following error. Earlier I was using categorical cross-entropy loss and metric as accuracy where it was compiling just fine. However, soon I realized that it is more appropriate to use mIOU or Jaccard index as metrics but am facing the following error. ( I think it is because though the "values" are the same one is a tensor? I might be wrong)

the error:

ValueError: in user code:

    File "/usr/local/lib/python3.7/dist-packages/keras/engine/training.py", line 1021, in train_function  *
        return step_function(self, iterator)
    File "<ipython-input-24-4f4d58e1a657>", line 11, in jacard_coef_loss  *
        return -jacard_coef(y_true, y_pred)  # -1 ultiplied as we want to minimize this value as loss function
    File "<ipython-input-24-4f4d58e1a657>", line 6, in jacard_coef  *
        intersection = K.sum(y_true_f * y_pred_f)

    ValueError: Dimensions must be equal, but are 1048576 and 4194304 for '{{node jacard_coef_loss/mul}} = Mul[T=DT_FLOAT](jacard_coef_loss/Reshape, jacard_coef_loss/Reshape_1)' with input shapes: [1048576], [4194304].

my data set shape :

Train Dataset: <BatchDataset element_spec=(TensorSpec(shape=(4, 512, 512, 3), dtype=tf.float32, name=None), TensorSpec(shape=(4, 512, 512, 1), dtype=tf.float32, name=None))>

Val Dataset: <BatchDataset element_spec=(TensorSpec(shape=(4, 512, 512, 3), dtype=tf.float32, name=None), TensorSpec(shape=(4, 512, 512, 1), dtype=tf.float32, name=None))>

Code for model creation- referred Keras deeplabv3+ sample

def convolution_block(
    block_input,
    num_filters=256,
    kernel_size=3,
    dilation_rate=1,
    padding="same",
    use_bias=False,
):
    x = layers.Conv2D(
        num_filters,
        kernel_size=kernel_size,
        dilation_rate=dilation_rate,
        padding="same",
        use_bias=use_bias,
        kernel_initializer=keras.initializers.HeNormal(),
    )(block_input)
    x = layers.BatchNormalization()(x)
    return tf.nn.relu(x)


def DilatedSpatialPyramidPooling(dspp_input):
    dims = dspp_input.shape
    x = layers.AveragePooling2D(pool_size=(dims[-3], dims[-2]))(dspp_input)
    x = convolution_block(x, kernel_size=1, use_bias=True)
    out_pool = layers.UpSampling2D(
        size=(dims[-3] // x.shape[1], dims[-2] // x.shape[2]), interpolation="bilinear",
    )(x)

    out_1 = convolution_block(dspp_input, kernel_size=1, dilation_rate=1)
    out_6 = convolution_block(dspp_input, kernel_size=3, dilation_rate=6)
    out_12 = convolution_block(dspp_input, kernel_size=3, dilation_rate=12)
    out_18 = convolution_block(dspp_input, kernel_size=3, dilation_rate=18)

    x = layers.Concatenate(axis=-1)([out_pool, out_1, out_6, out_12, out_18])
    output = convolution_block(x, kernel_size=1)
    return output

def DeeplabV3Plus(image_size, num_classes):
    model_input = keras.Input(shape=(image_size, image_size, 3))
    resnet50 = keras.applications.ResNet50(
        weights="imagenet", include_top=False, input_tensor=model_input
    )
    x = resnet50.get_layer("conv4_block6_2_relu").output
    x = DilatedSpatialPyramidPooling(x)

    input_a = layers.UpSampling2D(
        size=(image_size // 4 // x.shape[1], image_size // 4 // x.shape[2]),
        interpolation="bilinear",
    )(x)
    input_b = resnet50.get_layer("conv2_block3_2_relu").output
    input_b = convolution_block(input_b, num_filters=48, kernel_size=1)

    x = layers.Concatenate(axis=-1)([input_a, input_b])
    x = convolution_block(x)
    x = convolution_block(x)
    x = layers.UpSampling2D(
        size=(image_size // x.shape[1], image_size // x.shape[2]),
        interpolation="bilinear",
    )(x)
    model_output = layers.Conv2D(num_classes, kernel_size=(1, 1), padding="same")(x)
    return keras.Model(inputs=model_input, outputs=model_output)


model = DeeplabV3Plus(image_size=IMAGE_SIZE, num_classes=NUM_CLASSES)
model.summary()

the model:

---

 Layer (type)                   Output Shape         Param #     Connected to                     
==================================================================================================
 input_3 (InputLayer)           [(None, 512, 512, 3  0           []                               
                                )]                                                                
                                                                                                  
 conv1_pad (ZeroPadding2D)      (None, 518, 518, 3)  0           ['input_3[0][0]']                
                                                                                                  
 conv1_conv (Conv2D)            (None, 256, 256, 64  9472        ['conv1_pad[0][0]']              
                                )                                                                 
                                                                                                  
 conv1_bn (BatchNormalization)  (None, 256, 256, 64  256         ['conv1_conv[0][0]']             
                                )                                                                 
                                                                                                  
 conv1_relu (Activation)        (None, 256, 256, 64  0           ['conv1_bn[0][0]']               
                                )                                                                 
                                                                                                  
 pool1_pad (ZeroPadding2D)      (None, 258, 258, 64  0           ['conv1_relu[0][0]']             
                                )                                                                 
                                                                                                  
 pool1_pool (MaxPooling2D)      (None, 128, 128, 64  0           ['pool1_pad[0][0]']              
                                )                                                                 
                                                                                                  
 conv2_block1_1_conv (Conv2D)   (None, 128, 128, 64  4160        ['pool1_pool[0][0]']             
                                )                                                                 
                                                                                                  
 .                                                            
 .
                                                                                            
 up_sampling2d_7 (UpSampling2D)  (None, 128, 128, 25  0          ['tf.nn.relu_23[0][0]']          
                                6)                                                                
                                                                                                  
 tf.nn.relu_24 (TFOpLambda)     (None, 128, 128, 48  0           ['batch_normalization_24[0][0]'] 
                                )                                                                 
                                                                                                  
 concatenate_5 (Concatenate)    (None, 128, 128, 30  0           ['up_sampling2d_7[0][0]',        
                                4)                                'tf.nn.relu_24[0][0]']          
                                                                                                  
 conv2d_27 (Conv2D)             (None, 128, 128, 25  700416      ['concatenate_5[0][0]']          
                                6)                                                                
                                                                                                  
 batch_normalization_25 (BatchN  (None, 128, 128, 25  1024       ['conv2d_27[0][0]']              
 ormalization)                  6)                                                                
                                                                                                  
 tf.nn.relu_25 (TFOpLambda)     (None, 128, 128, 25  0           ['batch_normalization_25[0][0]'] 
                                6)                                                                
                                                                                                  
 conv2d_28 (Conv2D)             (None, 128, 128, 25  589824      ['tf.nn.relu_25[0][0]']          
                                6)                                                                
                                                                                                  
 batch_normalization_26 (BatchN  (None, 128, 128, 25  1024       ['conv2d_28[0][0]']              
 ormalization)                  6)                                                                
                                                                                                  
 tf.nn.relu_26 (TFOpLambda)     (None, 128, 128, 25  0           ['batch_normalization_26[0][0]'] 
                                6)                                                                
                                                                                                 
 up_sampling2d_8 (UpSampling2D)  (None, 512, 512, 25  0          ['tf.nn.relu_26[0][0]']          
                                6)                                                                
                                                                                                
 conv2d_29 (Conv2D)             (None, 512, 512, 4)  1028        ['up_sampling2d_8[0][0]']        

================================================================================================== Total params: 11,853,124 Trainable params: 11,820,388 Non-trainable params: 32,736 }

compiler and fitter which error occurs when I wanna fit the model

from keras import backend as K

def jacard_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 (intersection + 1.0) / (K.sum(y_true_f) + K.sum(y_pred_f) - intersection + 1.0)


def jacard_coef_loss(y_true, y_pred):
    return -jacard_coef(y_true, y_pred) 

model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.001),loss = jacard_coef_loss, metrics = jacard_coef)

history = model.fit(train_dataset, validation_data=val_dataset, epochs=10)