In my machine learning problem I am using a CNN to classify images.

Since my dataset is imbalanced I want to perform classification probability threshold tuning so I can find the optimal balance between sensitivity and specificity using G-mean over the returned values of the ROC-curves (explanation link).

I have split the dataset into train-test. I want to do Cross Validation on the training set only so I can hypertune the network parameters (mainly batch size, number of epochs and optimizer). For finding the best model I am thinking of using the roc_auc_score since its value is not dependent on the classification threshold applied.

I think that is probably a bad practice to use the test set to do this threshold determination directly (although I have seen it done).

My question comes down to this: After finding the best model how can I find the optimal threshold value? I am thinking of the following possibilities:

  • A: Assume that the best set of model/parameters doesn't overfit on the training data (since it's the best performing model on the gridsearch) and retrain on the full training set, obtain the classification probabilities for the training set, get optimal threshold.
  • B: Instead of dividing the original dataset into train-test, create a third set (so train-validation-test split). Find the best model on training set (through gridSearchCV), retrain on the training set and predict the validation set, optimize the threshold on this prediction.
  • C: Use part of the training data to train( like 70%), another part to validate. Find the optimal threshold.
  • D: With the best model perform CV again. Log the CV results for each fold. For a set of threshold let's say [0.1, 0.2, ..., 0.9] check which one across the folds has better average score (for example f1-macro).

On all the possibilities, I would then use the full training data (+ validation on the case of B and C) to train the model, get the test classification probabilities, use the previously determined threshold to the prediction classes, compute classification metrics.

I think that in A, if my best model overfits, the threshold will be wrong (although I can always compare it to the 0.5 threshold and see if it improves or not).

In B, I am giving up samples that could be helpful on the gridsearchCV.

In C, like in A I am also retraining on data that was use for model selection. The difference to A is that the prediction is not done on data that has been fit.

I don't really see any disadvantage in D. Although I didn't find that suggestion anywhere, which I find a little weird. I guess I don't really need to perform CV again if I just log the best threshold for all the gridsearched models.

What do you think is the best approach? Is there a better way to do this?



Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.