Boels Maxence
Boels Maxence

Reputation: 359

Output layer for Binary Classification in Keras

I need to classify images as either cancerous or not cancerous.

For this, I built a classical CNN but I am hesitating between labeling my dataset with either two-column vector like this:

cancerous: [0, 1]
not cancerous: [1, 0]

and using a softmax activation function with 2 output neurons.

model.add(Dense(2, activation='softmax'))

OR

cancerous: [1]
not cancerous: [0]

and using a sigmoid activation function with one output neuron.

model.add(Dense(1, activation='sigmoid'))

Which model is better given that I need to use the probability of having cancer as final metric for the patient and also for plotting the ROC curve?

Upvotes: 6

Views: 9776

Answers (1)

Ashwin Geet D'Sa
Ashwin Geet D'Sa

Reputation: 7369

The general tendency is to use multiple output nodes with sigmoid curve for multi-label classification. Often, a softmax is used for multiclass classification, where softmax predicts the probabilities of each output and we choose class with highest probability. For binary classification, we can choose a single neuron output passed through sigmoid, and then set a threshold to choose the class, or use two neuron output and then perform a softmax. In either of the cases, thresholding is possible.It is rather easy to plot a ROC curve with single neuron output, as you'll have to threshold over one value. So, you can easily go with model.add(Dense(1, activation='sigmoid'))

Upvotes: 4

Related Questions