Batching Soft IoU for Training Semantic Segmentation Networks

摘要

The majority of semantic segmentation networks generally employ cross-entropy as a loss function and intersection-over-union (IoU) as the evaluation metric for network performance. Employing IoU as a loss function can solve the mismatch issue between the loss function and the evaluation metric. We propose a Soft IoU training strategy based on mini-batch (mini-batch Soft IoU). Our work has two primary contributions: The first is to extend the IoU loss function to a multi-class segmentation network. The second is to collect various categories of the training samples in every mini-batch, which will ensure that the number of categories equals to the batch size at least. Our method breaks the randomness of the original mini-batch gradient descent (GD) strategy, advancing training samples in the mini-batch much more consistent with the distribution characteristics of the overall data. It solves the instability of IoU loss function. In addition, the experimental results on the PASCAL VOC2012 dataset reveal that our method effectively improves the segmentation accuracy of the network and attains significant improvements beyond state-of-the-art IoU loss function methods.