I am training a U-Net for semantic segmentation of large medical images (4096x4096px). The two classes are "too" unbalanced. The white pixels are just about 0.1% (or less) of the whole image. The Dice Coeff loss function seems to not be working since it predicts always black pixels.
- Is there any specialized loss function for such unbalanced data? I can not find anything that works.
- Is the U-Net arcitecture suitable for such segmentation tasks?
I have tried to train with the following setup:
Epochs: 50 Batch size: 4 Learning rate: 1e-05 Training size: 451 Validation size: 23 Checkpoints: True Device: cuda Images scaling: 0.25
and also with batch size of 1 and learning rate of 10^4. I would appreciate some help.
Cheers
Some feedback:
1. UNet is outdated. Try Mask-RCNN or YOLO.
2. Training size is little less. I don't know exactly what medical problem you're addressing, but can you try some augmentation? Say, mirror the image or translate/rotate.
3. You may try increasing the batch size to say 4.
4. for loss function, you may actually alter dice score to weigh positive class (white class) more.
Hi!
I'd object to the notion of U-Net being "outdated". It's a perfectly good solution to a lot of tasks.
In my experience it pays to add a couple more epochs in such cases. Also combining dice loss and binary cross entropy by using their (weighted) sum is much more stable than pure Dice loss in my experience.
- Badges
- Science topic