Model is not learning

I am trying to generate image segmentation for cerebral aneurysms. I have used UNET resnet50-unet and dense-net. UNET. after 5 epochs the valid loss increases or doesn’t change. I have done data augmentation. The images are Maximum intensity projection of brain MRIs. The loss starts at 0.99 and the lowest it goes is 0.95. The Jaccard score is 0.04 (highest) masks generated by the model are vague. Kernel size has been changed from 3 to 5 , learning rate is 0.001, optimiser Adam, loss function dice loss.

The code for model is:
from tensorflow.keras.layers import Conv2D, BatchNormalization, Activation, MaxPool2D, Conv2DTranspose, Concatenate, Input
from tensorflow.keras.models import Model
from tensorflow.keras.applications import ResNet50

def conv_block(input, num_filters):
x = Conv2D(num_filters, 5, padding=“same”)(input)
x = BatchNormalization()(x)
x = Activation(“relu”)(x)

x = Conv2D(num_filters, 5, padding="same")(x)
x = BatchNormalization()(x)
x = Activation("relu")(x)

return x

def decoder_block(input, skip_features, num_filters):
x = Conv2DTranspose(num_filters, (2, 2), strides=2, padding=“same”)(input)
x = Concatenate()([x, skip_features])
x = conv_block(x, num_filters)
return x

def build_resnet50_unet(input_shape):
“”" Input “”"
inputs = Input(input_shape)

""" Pre-trained ResNet50 Model """
resnet50 = ResNet50(include_top=False, weights=None, input_tensor=inputs)

""" Encoder """
s1 = resnet50.layers[0].output          ## (512 x 512)
s2 = resnet50.get_layer("conv1_relu").output        ## (256 x 256)
s3 = resnet50.get_layer("conv2_block3_out").output  ## (128 x 128)
s4 = resnet50.get_layer("conv3_block4_out").output  ## (64 x 64)

""" Bridge """
b1 = resnet50.get_layer("conv4_block6_out").output  ## (32 x 32)

""" Decoder """
d1 = decoder_block(b1, s4, 512)                     ## (64 x 64)
d2 = decoder_block(d1, s3, 256)                     ## (128 x 128)
d3 = decoder_block(d2, s2, 128)                     ## (256 x 256)
d4 = decoder_block(d3, s1, 64)                      ## (512 x 512)

""" Output """
outputs = Conv2D(1, 1, padding="same", activation="sigmoid")(d4)

model = Model(inputs, outputs, name="ResNet152_U-Net")
return model

if name == “main”:
input_shape = (512, 512, 3)
model = build_resnet50_unet(input_shape)
model.summary()

The code for training is:

import os
import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K
os.environ[“TF_CPP_MIN_LOG_LEVEL”] = “2” #set to 1 for warnings and errors
import numpy as np
import cv2
import keras
import keras.utils

from glob import glob
from sklearn.utils import shuffle
import tensorflow as tf
from tensorflow.keras.callbacks import ModelCheckpoint, CSVLogger, ReduceLROnPlateau, EarlyStopping, TensorBoard
from tensorflow.keras.optimizers import Adam
from tensorflow.keras.metrics import Recall, Precision
H = 512
W = 512

def create_dir(path):
“”" Create a directory. “”"
if not os.path.exists(path):
os.makedirs(path)

def shuffling(x, y):
x, y = shuffle(x, y, random_state=42)
return x, y

def load_data(path):
x = sorted(glob(os.path.join(path, “image”, “.png")))
y = sorted(glob(os.path.join(path, “mask”, "
.png”)))
return x, y

def read_image(path):
path = path.decode()
x = cv2.imread(path, cv2.IMREAD_COLOR)
x = x/255.0
x = x.astype(np.float32)
return x

def read_mask(path):

path = path.decode()
x = cv2.imread(path, cv2.IMREAD_GRAYSCALE)
x = x/255.0

x = x > 0.5
x = x.astype(np.float32)
x = np.expand_dims(x, axis=-1)
return x

def tf_parse(x, y):
def _parse(x, y):
x = read_image(x)
y = read_mask(y)
return x, y

x, y = tf.numpy_function(_parse, [x, y], [tf.float32, tf.float32])
x.set_shape([H, W, 3])
y.set_shape([H, W, 1])
return x, y

def tf_dataset(x, y, batch=8):
dataset = tf.data.Dataset.from_tensor_slices((x, y))
dataset = dataset.map(tf_parse)
dataset = dataset.batch(batch)
dataset = dataset.prefetch(10)
return dataset

if name == “main”:
“”" Seeding “”"
np.random.seed(42)
tf.random.set_seed(42)

""" Directory for storing files """
create_dir("files")

""" Hyperparameters """
batch_size = 10
lr = 0.001
num_epochs = 200
model_path = os.path.join("files", "model.h5")
csv_path = os.path.join("files", "data.csv")

""" Dataset """
train_path = os.path.join("/content/drive/MyDrive/Data_brain/train/")
valid_path = os.path.join("/content/drive/MyDrive/Data_brain/test/")

train_x, train_y = load_data(train_path)
train_x, train_y = shuffling(train_x, train_y)
valid_x, valid_y = load_data(valid_path)

print(f"Train: {len(train_x)} - {len(train_y)}")
print(f"Valid: {len(valid_x)} - {len(valid_y)}")

train_dataset = tf_dataset(train_x, train_y, batch=batch_size)
valid_dataset = tf_dataset(valid_x, valid_y, batch=batch_size)

""" Model """
model = build_resnet50_unet(input_shape)
metrics = [dice_coef, Recall(), Precision()]
model.compile(loss=dice_loss, optimizer=Adam(lr), metrics=metrics)

callbacks = [
    ModelCheckpoint(model_path, verbose=1, save_best_only=True),
    #ReduceLROnPlateau(monitor='val_loss', factor=0.1, patience=10, min_lr=1e-7, verbose=1),
    CSVLogger(csv_path),
    TensorBoard(),
    #EarlyStopping(monitor='val_loss', patience=50, restore_best_weights=False),
]

model.fit(
    train_dataset,
    epochs=num_epochs,
    validation_data=valid_dataset,
    callbacks=callbacks,
    shuffle=False)

code for dice loss:
import numpy as np
import tensorflow as tf
from tensorflow.keras import backend as K

def iou(y_true, y_pred):
def f(y_true, y_pred):
intersection = (y_true * y_pred).sum()
union = y_true.sum() + y_pred.sum() - intersection
x = (intersection + 1e-15) / (union + 1e-15)
x = x.astype(np.float32)
return x
return tf.numpy_function(f, [y_true, y_pred], tf.float32)

smooth = 1e-15
def dice_coef(y_true, y_pred):
y_true = tf.keras.layers.Flatten()(y_true)
y_pred = tf.keras.layers.Flatten()(y_pred)
intersection = tf.reduce_sum(y_true * y_pred)
return (2. * intersection + smooth) / (tf.reduce_sum(y_true) + tf.reduce_sum(y_pred) + smooth)

def dice_loss(y_true, y_pred):
return 1.0 - dice_coef(y_true, y_pred)

My guess is that custom loss function is the culprit. Try to train the model with a crossentropy loss.
Also you can read this discussion of dice loss: python - Implementing Multiclass Dice Loss Function - Stack Overflow

1 Like