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:
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:
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.