Keras U-net implementation for IRIS segmentation with weighted loss function

Question

I'm trying to segment the human IRIS images with U-net using the weighted loss function. I'm referring to this bloglink Below is the model definition

def my_loss(target, output):
    return - tf.reduce_sum(target * output,
                           len(output.get_shape()) - 1)
# Standard Unet model from blog post
_epsilon = tf.convert_to_tensor(K.epsilon(), np.float32)

def make_weighted_loss_unet(input_shape, n_classes):
    ip = L.Input(shape=input_shape)
    weight_ip = L.Input(shape=input_shape[:2] + (n_classes,))    

    conv1 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(ip)
    conv1 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv1)
    conv1 = L.Dropout(0.1)(conv1)
    mpool1 = L.MaxPool2D()(conv1)

    conv2 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool1)
    conv2 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv2)
    conv2 = L.Dropout(0.2)(conv2)
    mpool2 = L.MaxPool2D()(conv2)

    conv3 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool2)
    conv3 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv3)
    conv3 = L.Dropout(0.3)(conv3)
    mpool3 = L.MaxPool2D()(conv3)

    conv4 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool3)
    conv4 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv4)
    conv4 = L.Dropout(0.4)(conv4)
    mpool4 = L.MaxPool2D()(conv4)

    conv5 = L.Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(mpool4)
    conv5 = L.Conv2D(1024, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv5)
    conv5 = L.Dropout(0.5)(conv5)

    up6 = L.Conv2DTranspose(512, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv5)
    conv6 = L.Concatenate()([up6, conv4])
    conv6 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)
    conv6 = L.Conv2D(512, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv6)
    conv6 = L.Dropout(0.4)(conv6)

    up7 = L.Conv2DTranspose(256, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv6)
    conv7 = L.Concatenate()([up7, conv3])
    conv7 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)
    conv7 = L.Conv2D(256, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv7)
    conv7 = L.Dropout(0.3)(conv7)

    up8 = L.Conv2DTranspose(128, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv7)
    conv8 = L.Concatenate()([up8, conv2])
    conv8 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)
    conv8 = L.Conv2D(128, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv8)
    conv8 = L.Dropout(0.2)(conv8)

    up9 = L.Conv2DTranspose(64, 2, strides=2, kernel_initializer='he_normal', padding='same')(conv8)
    conv9 = L.Concatenate()([up9, conv1])
    conv9 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
    conv9 = L.Conv2D(64, 3, activation='relu', padding='same', kernel_initializer='he_normal')(conv9)
    conv9 = L.Dropout(0.1)(conv9)

  
    c10 = L.Conv2D(n_classes, 1, activation='sigmoid', kernel_initializer='he_normal')(conv9)

    #Mimic crossentropy loss
    c11 = L.Lambda(lambda x: x / tf.reduce_sum(x, len(x.get_shape()) - 1, True))(c10)
    c11 = L.Lambda(lambda x: tf.clip_by_value(x, _epsilon, 1. - _epsilon))(c11)
    c11 = L.Lambda(lambda x: K.log(x))(c11)
    weighted_sm = L.multiply([c11, weight_ip])

    model = Model(inputs=[ip, weight_ip], outputs=[weighted_sm])
    
    return model

This is how I'm training

model = make_weighted_loss_unet((256, 256, 3), 1) # shape of input, number of classes
model.compile(optimizer='adam',loss=my_loss, metrics=['acc'])
model.fit([X_train, wmap], y_train, validation_split=0.1, epochs=100)

Here

X_train: list of input images (imgs, 256, 256, 3)
wmap: list of weight maps (imgs, 256,256,1)
y_train:list of masks(imgs,256,256,256,1)

Training output:

Epoch 1/100
2/2 [==============================] - 39s 17s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887
Epoch 2/100
2/2 [==============================] - 40s 17s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887
Epoch 3/100
2/2 [==============================] - 41s 18s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887
Epoch 4/100
2/2 [==============================] - 39s 17s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887
Epoch 5/100
2/2 [==============================] - 39s 17s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887
Epoch 6/100
2/2 [==============================] - 40s 18s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887
Epoch 7/100
2/2 [==============================] - 40s 18s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887
Epoch 8/100
2/2 [==============================] - 41s 18s/step - loss: nan - acc: 0.7119 - val_loss: nan - val_acc: 0.6887

The model accuracy is not improving in fact stays unchanged, and losses are shown as nan for the "sigmoid" activation function. If I use "softmax" still the accuracy is unchanged, and the loss remains the same.

Did anyone faced the same issue before, or has anyone who has worked on this previously? Please suggest. Thanks in advance, Naveen

Keras U-net implementation for IRIS segmentation with weighted loss function

0 Answers