Keras: how to deal with unbalanced classes in sematic segmentation task? - deep-learning

Question

I have pretrained VGG16 based FCN-32s like model, defined like:
def pop_layer(model):
if not model.outputs:
raise Exception('Sequential model cannot be popped: model is empty.')
model.layers.pop()
if not model.layers:
model.outputs = []
model.inbound_nodes = []
model.outbound_nodes = []
else:
model.layers[-1].outbound_nodes = []
model.outputs = [model.layers[-1].output]
model.built = False
def get_model():
#Fully convolutional part of VGG16
model = VGG16(include_top=False, weights='imagenet')
#Remove last max pooling layer
pop_layer(model)
#Freeze pretrained layers
for layer in model.layers:
layer.trainable = False
model = Model(inputs=model.inputs, outputs=model.outputs)
#print('len(model.layers)', len(model.layers)) #
#print(model.summary()) #
x = Conv2D(512, (3, 3), activation='relu', padding='same')(model.output)
x = Conv2DTranspose(NUMBER_OF_CLASSES, kernel_size=(32, 32), strides=(16, 16), activation='sigmoid', padding='same')(x)
head = Reshape((-1,NUMBER_OF_CLASSES))(x)
model = Model(inputs=model.inputs, outputs=head)
model.compile(optimizer=Adadelta(), loss='binary_crossentropy')
print('len(model.layers)', len(model.layers)) #
print(model.summary()) #
return model
Model summary:
len(model.layers) 21
_________________________________________________________________
Layer (type) Output Shape Param #
=================================================================
input_1 (InputLayer) (None, None, None, 3) 0
_________________________________________________________________
block1_conv1 (Conv2D) (None, None, None, 64) 1792
_________________________________________________________________
block1_conv2 (Conv2D) (None, None, None, 64) 36928
_________________________________________________________________
block1_pool (MaxPooling2D) (None, None, None, 64) 0
_________________________________________________________________
block2_conv1 (Conv2D) (None, None, None, 128) 73856
_________________________________________________________________
block2_conv2 (Conv2D) (None, None, None, 128) 147584
_________________________________________________________________
block2_pool (MaxPooling2D) (None, None, None, 128) 0
_________________________________________________________________
block3_conv1 (Conv2D) (None, None, None, 256) 295168
_________________________________________________________________
block3_conv2 (Conv2D) (None, None, None, 256) 590080
_________________________________________________________________
block3_conv3 (Conv2D) (None, None, None, 256) 590080
_________________________________________________________________
block3_pool (MaxPooling2D) (None, None, None, 256) 0
_________________________________________________________________
block4_conv1 (Conv2D) (None, None, None, 512) 1180160
_________________________________________________________________
block4_conv2 (Conv2D) (None, None, None, 512) 2359808
_________________________________________________________________
block4_conv3 (Conv2D) (None, None, None, 512) 2359808
_________________________________________________________________
block4_pool (MaxPooling2D) (None, None, None, 512) 0
_________________________________________________________________
block5_conv1 (Conv2D) (None, None, None, 512) 2359808
_________________________________________________________________
block5_conv2 (Conv2D) (None, None, None, 512) 2359808
_________________________________________________________________
block5_conv3 (Conv2D) (None, None, None, 512) 2359808
_________________________________________________________________
conv2d_1 (Conv2D) (None, None, None, 512) 2359808
_________________________________________________________________
conv2d_transpose_1 (Conv2DTr (None, None, None, 3) 1572867
_________________________________________________________________
reshape_1 (Reshape) (None, None, 3) 0
=================================================================
Total params: 18,647,363
Trainable params: 3,932,675
Non-trainable params: 14,714,688
_________________________________________________________________
None
But when I train model it only predict most dominant class, my dataset is unbalanced:
Pixel area per class ratio:
class1 : 62.93 %
class2 : 25.46 %
class3 : 11.61 %
So my questions are: is my model definition ok? How to deal with class inbalanced? maybe batch should be constructed in some special way?

Answer 1

It seems like your loss doesn't fit your problem. You use binary cross entropy loss here:
model.compile(optimizer=Adadelta(), loss='binary_crossentropy')
But you have more than two classes. So I would suggest to use categorical_crossentropy loss (appears in the list of losses here. Read on the bottom of the page how to prepare your data in order to use this loss).
There are additional types of losses which may fit better inbalance classes situation. You may try to use dice loss, which is a differential approximation of IoU (intersection over union).
This loss is described on page 6, section 3 here.

Answer 2

A common approach is to use class weighing for the loss function, so you can penalize the effect of the prominent class.
weight_class = 1/ln(c + class_probability)
where c is a constant and mostly used value 1.03.