I have implemented two CNN architectures to perform segmentations on medical images: the classic UNet and a modified version called the Attention UNet. I have been training the models on roughly 50,000 images with 10-fold CV, 10 epochs, and batch size 16. The loss functions seem to get fairly small <0.1 throughout the trainings yet my issue lies in the predicted outputs. Many times I would train a model and receive "nice-looking" heat maps as the predicted output like the one below:

enter image description here

These predictions tend to produce the highest test set Dice scores after evaluating the appropriate cutoff pixel threshold.

Yet, currently, upon training the very same model (often with slight adjustments), I get these predictions which produce worse scores:

enter image description here

I am fairly new at CNNs and segmentation tasks and, instead of dumping my entire codebase to track down the issue, I would like to know if there is any general principle in DL specifically with CNNs that might be causing this issue. I have tried everything from tracking my changes to the model back to when I received the heat maps as well as implementing better training techniques that prevent overfitting, to no avail. Any insight or hint in the right direction is much appreciated.



Your Answer

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