Number of units in Final softmax layer in VGGNet16

Question

I am trying to implement and train neural network model VGGNet from scratch, on my own data. I am reproducing all the layers of the model. I am having a confusion about the last, fully connected softmax layer.

In the research paper by Simonyan and Zisserman, the last layer of VGGNet has 1000 neurons with sotmax function. As I understand this, this layer can classify data into 1000 classes. However, my data has only 4 classes. So should I add a layer with only 4 units instead, or should I add a 1000 unit layer, and add a 4 unit layer on top of it?

Answer 1

You should do the first one: add a layer with only 4 units instead of the 1000 unit layer.

You can think of the first $N-1$ layers as a feature extractor, converting the high-dimensional image to a (relatively) low-dimensional dense vector, and the $N$th layer as a linear classifier over that dense vector. The goal during training is then to learn a feature extractor such that the output is linearly separable with respect to the classes.

Because you changed the number of classes, you need to change the size of the output layer.