# Is it reasonable to use VGG16 with a new fully-connected layer for binary image segmentation?

I am working on binary image segmentation of traffic signs (of which I have RGB images of size 224x224 and accompanying grayscale masks) where I want to classify each pixel as either part of a traffic sign (1) or not (0), i.e. foreground (1) or background (0).

I wanted to start simple and try the following approach:

1. Take tf.keras.applications.vgg16.VGG16 pre-trained on ImageNet and freeze all layers (i.e. don't train them)
2. Pop off the last Dense layer of 1000 units (one for each of the original 1000 image classes)
3. Install my own Dense layer of 50176 units (one for each of the 224*224=50176 pixels)
4. Follow it by a Reshape to (224, 224, 1).
5. Train the network (basically the last layer I installed) on my dataset

In code my model looks like:

def vgg16(img_height, img_width, output_activation, loss, optimizer):
# Freeze VGG16's layers
vgg16 = tf.keras.applications.vgg16.VGG16(weights='imagenet')
for layer in vgg16.layers:
layer.trainable = False

# Stitch VGG16 to our own fully-connected layer for pixel-wise classification
x = vgg16.get_layer('fc2').output
x = tf.keras.layers.Dense(units=img_height*img_width, activation=output_activation)(x)
x = tf.keras.layers.Reshape((img_height, img_width, 1))(x)
model = tf.keras.models.Model(inputs=vgg16.input, outputs=x)
model.compile(loss=loss, optimizer=optimizer, metrics=['accuracy'])
return model


The model.summary() looks like:

_________________________________________________________________
Layer (type)                 Output Shape              Param #
=================================================================
input_1 (InputLayer)         (None, 224, 224, 3)       0
_________________________________________________________________
block1_conv1 (Conv2D)        (None, 224, 224, 64)      1792
_________________________________________________________________
block1_conv2 (Conv2D)        (None, 224, 224, 64)      36928
_________________________________________________________________
block1_pool (MaxPooling2D)   (None, 112, 112, 64)      0
_________________________________________________________________
block2_conv1 (Conv2D)        (None, 112, 112, 128)     73856
_________________________________________________________________
block2_conv2 (Conv2D)        (None, 112, 112, 128)     147584
_________________________________________________________________
block2_pool (MaxPooling2D)   (None, 56, 56, 128)       0
_________________________________________________________________
block3_conv1 (Conv2D)        (None, 56, 56, 256)       295168
_________________________________________________________________
block3_conv2 (Conv2D)        (None, 56, 56, 256)       590080
_________________________________________________________________
block3_conv3 (Conv2D)        (None, 56, 56, 256)       590080
_________________________________________________________________
block3_pool (MaxPooling2D)   (None, 28, 28, 256)       0
_________________________________________________________________
block4_conv1 (Conv2D)        (None, 28, 28, 512)       1180160
_________________________________________________________________
block4_conv2 (Conv2D)        (None, 28, 28, 512)       2359808
_________________________________________________________________
block4_conv3 (Conv2D)        (None, 28, 28, 512)       2359808
_________________________________________________________________
block4_pool (MaxPooling2D)   (None, 14, 14, 512)       0
_________________________________________________________________
block5_conv1 (Conv2D)        (None, 14, 14, 512)       2359808
_________________________________________________________________
block5_conv2 (Conv2D)        (None, 14, 14, 512)       2359808
_________________________________________________________________
block5_conv3 (Conv2D)        (None, 14, 14, 512)       2359808
_________________________________________________________________
block5_pool (MaxPooling2D)   (None, 7, 7, 512)         0
_________________________________________________________________
flatten (Flatten)            (None, 25088)             0
_________________________________________________________________
fc1 (Dense)                  (None, 4096)              102764544
_________________________________________________________________
fc2 (Dense)                  (None, 4096)              16781312
_________________________________________________________________
dense (Dense)                (None, 50176)             205571072
_________________________________________________________________
reshape (Reshape)            (None, 224, 224, 1)       0
=================================================================
Total params: 339,831,616
Trainable params: 205,571,072
Non-trainable params: 134,260,544
_________________________________________________________________


Accuracy (on the training data) is reported to be around 80-90% which seemed promising to me even though I was still just preparing the basic training setup. The problem at this point is that the model's predictions (also on the training data, which should be good according to training accuracy) are garbage. I am still looking for silly programming or engineering mistakes, but I am starting to question my entire approach so I'll put that aside for now.

Is my approach reasonable from a neural network design point of view? In short, is it reasonable to take VGG16, pop off the last fully-connected layer and install my own with 50176 units (one for each pixel) and reshape the output to (224, 224, 1)?

This is perfectly reasonable and a common way of transferring knowledge from VGG16 to your own network.

As the C layers are basically doing image segmentation and feature selection for the purpose of VGG16, what you need also to ask yourself is if what you want to do is "similar" to what VGG16 does. If the answer is true, then the features VGG16 extracts should be a good starting point for your training.

Of course, design your network so that the number of parameters is reasonable compared to the amount of samples you have.

• The main reason I'm asking is because most of the segmentation networks I have seen follow an encoder/decoder architecture, the reason which is briefly explained in the first couple of paragraphs in jeremyjordan.me/semantic-segmentation. With that in mind, is there any inherent reason why my network may not work? I am not really after state-of-the-art performance, I am researching something else entirely. – fabiomaia Dec 26 '18 at 0:11
• Indeed, if you want to try this kind of architectures, the last bits won't be provided by VGG16. Still, you can start with VGG16 as the input and train the back, that's far less parameters to train than a full network and also a far better starting point. – Matthieu Brucher Dec 26 '18 at 9:45
• Yes, I wanted to stick to a transfer learning approach in order to save time, but actually the fully-connected layer with 50176 units I installed has 250 million parameters which is twice as many as the parameters in the previous layers I froze. Is the number of parameters perhaps the big advantage of those encoder/decoder architectures? – fabiomaia Dec 26 '18 at 19:21
• The main advantage is the reduction of the dimensionality of the problems that you are solving. The second can be seen as the number of parameters, but I don't think it's in the same range. The reduction of number of parameters is due to the convolutional outputs, and these outputs make sense for an image. You may want to use less dense intermediate nodes and add these transposed convolution layers up to the output. – Matthieu Brucher Dec 26 '18 at 20:12