Is there any example of using weighted loss for pixel-wise segmentation/classification tasks?

Question

I am doing FCN32 semantic segmentation on my data. I ran the algorithm to fine-tune for my data (grayscale images with only one channel), till 80,000 iterations; however, the loss and accuracy are fluctuating and the output image completely black. Even, the loss is so high after 80,000 iterations. I thought the classifier cannot do training well on my data. So, I am going to train from scratch. On the other hand, my data has imbalanced class members. The background pixels are more than the other four classes. Some researchers are suggesting using weighted loss. Does anyone have any idea? Am I doing the right way? How can I add this weighted loss to train_val.prototxt?

I will be thankful if you know any resources/examples related to training with weighted loss, please share with me here.

Thanks again

Answer 1

You can tackle class imbalance using "InfogainLoss". This loss can be viewed as an extension to "SoftmaxWithLoss" that enables you to "pay" different loss value per label.
If you want to use "InfogainLoss" for pixel-wise predictions, you might need to use BVLC/caffe PR#3855.