# Unet Overfitting for binary segmentation of fake images

I am working on a project where I am trying to detect and localize forgeries in images. I am using the CASIA v2 dataset and using Unet model for the task. I have the binary masks of all the images in the CASIA v2 dataset. The metric I am using for the model are F1 score.

The issue with the model is that it is highly overfitting, the validation loss plateaus up.

Batch size is 128 and Learning rate is 0.000001. Image size is 128 x 128.

Updated graph for batch size 16 with the changes mentioned by @spb is as follows:

I have also tried using Learning rate scheduler to decrease the learning rate(starting with high learning rate) on plateaus but that didn't help much.

I am also using the package Albumentations for data augmentation of both the images and its masks. I load the images and the masks and then apply the augmentations and save the augmented images and masks in a separate arrays and finally extend the original images and masks with the augmented images and masks. So technically I have original plus the augmented images and masks that I use for training the model. The augmentations I am using are:

Augment = A.Compose([
A.VerticalFlip(p=0.5),
A.RandomRotate90(p=0.5),
A.HorizontalFlip(p = 0.5)
])


I have split the dataset into 70% Training, 20% Validation and 10% for testing. Here is a snippet of my model. Updated Code below

def conv2d_block(input_tensor, n_filters, kernel_size = 3, batchnorm = True):
"""Function to add 2 convolutional layers with the parameters passed to it"""
# first layer
x = Conv2D(filters = n_filters, kernel_size = (kernel_size, kernel_size),\
kernel_initializer = 'he_normal', padding = 'same')(input_tensor)
if batchnorm:
x = BatchNormalization()(x)
x = Activation('relu')(x)

# second layer
x = Conv2D(filters = n_filters, kernel_size = (kernel_size, kernel_size),\
kernel_initializer = 'he_normal', padding = 'same')(input_tensor)
if batchnorm:
x = BatchNormalization()(x)
x = Activation('relu')(x)

return x

def get_unet(input_img, n_filters = 16, dropout = 0.1, batchnorm = True):
"""Function to define the UNET Model"""
# Contracting Path
c1 = conv2d_block(input_img, n_filters * 1, kernel_size = 3, batchnorm = batchnorm)
p1 = MaxPooling2D((2, 2))(c1)
#p1 = Dropout(dropout)(p1)

c2 = conv2d_block(p1, n_filters * 2, kernel_size = 3, batchnorm = batchnorm)
p2 = MaxPooling2D((2, 2))(c2)
#p2 = Dropout(dropout)(p2)

c3 = conv2d_block(p2, n_filters * 4, kernel_size = 3, batchnorm = batchnorm)
p3 = MaxPooling2D((2, 2))(c3)
#p3 = Dropout(dropout)(p3)

c4 = conv2d_block(p3, n_filters * 8, kernel_size = 3, batchnorm = batchnorm)
p4 = MaxPooling2D((2, 2))(c4)
#p4 = Dropout(dropout)(p4)

c5 = conv2d_block(p4, n_filters * 16, kernel_size = 3, batchnorm = batchnorm)
p5 = MaxPooling2D((2, 2))(c5)
#p5 = Dropout(dropout)(p5)

c6 = conv2d_block(p5, n_filters = n_filters * 32, kernel_size = 3, batchnorm = batchnorm)

# Expansive Path
u7 = Conv2DTranspose(n_filters * 16, (3, 3), strides = (2, 2), padding = 'same')(c6)
u7 = concatenate([u7, c5])
u7 = Dropout(dropout)(u7)
c7 = conv2d_block(u7, n_filters * 16, kernel_size = 3, batchnorm = batchnorm)

u8 = Conv2DTranspose(n_filters * 8, (3, 3), strides = (2, 2), padding = 'same')(c7)
u8 = concatenate([u8, c4])
u8 = Dropout(dropout)(u8)
c8 = conv2d_block(u8, n_filters * 8, kernel_size = 3, batchnorm = batchnorm)

u9 = Conv2DTranspose(n_filters * 4, (3, 3), strides = (2, 2), padding = 'same')(c8)
u9 = concatenate([u9, c3])
u9 = Dropout(dropout)(u9)
c9 = conv2d_block(u9, n_filters * 4, kernel_size = 3, batchnorm = batchnorm)

u10 = Conv2DTranspose(n_filters * 2, (3, 3), strides = (2, 2), padding = 'same')(c9)
u10 = concatenate([u10, c2])
u10 = Dropout(dropout)(u10)
c10 = conv2d_block(u10, n_filters * 2, kernel_size = 3, batchnorm = batchnorm)

u11 = Conv2DTranspose(n_filters * 1, (3, 3), strides = (2, 2), padding = 'same')(c10)
u11 = concatenate([u11, c1])
u11 = Dropout(dropout)(u11)
c11 = conv2d_block(u11, n_filters * 1, kernel_size = 3, batchnorm = batchnorm)

outputs = Conv2D(1, (1, 1), activation='sigmoid')(c11)
model = Model(inputs=[input_img], outputs=[outputs])
return model


Currently I am not using the dropout as it leads to higher validation loss plateaus in my case.

The F1 score and F1 loss I am calculating are as follows

def f1(y_true, y_pred):

y_pred = K.round(y_pred)
tp = K.sum(K.cast(y_true*y_pred, 'float'), axis=0)
tn = K.sum(K.cast((1-y_true)*(1-y_pred), 'float'), axis=0)
fp = K.sum(K.cast((1-y_true)*y_pred, 'float'), axis=0)
fn = K.sum(K.cast(y_true*(1-y_pred), 'float'), axis=0)

p = tp / (tp + fp + K.epsilon())
r = tp / (tp + fn + K.epsilon())

f1 = 2*p*r / (p+r+K.epsilon())
f1 = tf.where(tf.is_nan(f1), tf.zeros_like(f1), f1)
return K.mean(f1)

def f1_loss(y_true, y_pred):

tp = K.sum(K.cast(y_true*y_pred, 'float'), axis=0)
tn = K.sum(K.cast((1-y_true)*(1-y_pred), 'float'), axis=0)
fp = K.sum(K.cast((1-y_true)*y_pred, 'float'), axis=0)
fn = K.sum(K.cast(y_true*(1-y_pred), 'float'), axis=0)

p = tp / (tp + fp + K.epsilon())
r = tp / (tp + fn + K.epsilon())

f1 = 2*p*r / (p+r+K.epsilon())
f1 = tf.where(tf.is_nan(f1), tf.zeros_like(f1), f1)
return 1 - K.mean(f1)


I have also tried using other losses like focal_tversky but have a similar result.

What can be the issue and how can I solve it?

Is it

1. Issue with my data like presence of outliers
2. Model related issue
3. Batch size and Learning rate related issue
4. Or anything else?

Please your help in this regard is really appreciated as I really need to solve it soon.

I do not understand why you say that your model is overfitting. An overfit occurs when the validation loss start increasing after diminishing. Here it seems that your model has reaches its potential and cannot improve anymore. What I would recommend here is to make your model bigger: add filters, increase the depth. Also consider trying transfer learning; it is a common base to all tasks.

• Ok I will try working with a bigger model. Thanks for your inputs. I also updated my recent graph in the question above considering the inputs given by @spb and got these results. Do you think its overfitting or the model is basically reaching its potential? May 18 at 23:30
• Still, I don't see the shadow of an overfitting. However, I checked the dataset that you use, and it says to contain 4795 images. I think that training a model from scratch on a small dataset like this will not give you incredible results. Transfer learning can help a lot with that, combined with data augmentation. May 19 at 13:24
• Any recommendations which weights I shall use for transfer learning? May 19 at 14:36
• The U-Net architecture is basically an encoder and a decoder with skip connections between them. you can use any pre-trained deep ConvNet for the encoder part and then build your decoder accordingly with transposed convolutions; and don't forget to add skip connections - very important. Here you have the pre-trained models available in tensorflow that you can use for the encoder. May 19 at 14:46

Data augmentations is usually done on the fly during training, meaning before each you apply the random augmentation for the entire dataset, because of the randomness there will be different transformation of the same image in each epoch.

Shuffle the dataset before batching in each epoch, so that each epoch will not have minibatch of same images, which will reduce overfitting. Learning rate usually 1e-4 works fine for me.

Your UNet is not wide enough, why are you using only 16 filters in first conv block, original UNet paper had 64 filters in first conv block. Also you have only one convolution block in each layer, why? original unet has 2 conv blocks in each layer. I suggest you to try with unet given in here https://github.com/zhixuhao/unet/blob/master/model.py

Dice loss is usually prefered for segmentation, check code here

from keras import backend as K

def dice_score(y_true, y_pred, smooth=1e-7):

intersection = K.sum(K.abs(y_true * y_pred), axis=-1)
return (2. * intersection + smooth) / (K.sum(K.square(y_true),-1) + K.sum(K.square(y_pred),-1) + smooth)

def dice_loss(y_true, y_pred):
return 1-dice_score(y_true, y_pred)

• Thanks for your in detail answer. Just to clear some things, so random augmentations should be applied before each epoch on the entire dataset, right? Also I updated the code in the post above, I am using two convolution layers, conv2d_block is actually a function i made for each layer, i have put that code above as well for your reference. May 15 at 12:42
• Yes, apply random augmentations before each epoch on the entire dataset.
– spb
May 15 at 12:47
• What about number of filters in the first layer, did you try changing to 64? Also try removing dropouts in contracting path, it worked well for me.
– spb
May 15 at 12:48
• Yes, I am not using dropouts they are commented, so you used dropouts in the expanding path? I am going to try changing filters to 64 now and get back to you. May 15 at 12:53
• Yes, dropout(0.5) in expanding path reduced overfitting
– spb
May 15 at 12:56