IoU-Balanced loss functions for single-stage object detection

Single-stage object detectors have been widely applied in computer vision applications due to their high efficiency. However, the loss functions adopted by single-stage detectors hurt the localization accuracy seriously. Firstly, the cross-entropy loss for classification is independent of the localization task and drives all the positive examples to learn as high classification scores as possible regardless of localization accuracy. Thus, there exist many detections with high classification scores but low IoU or detections with low classification scores but high IoU. Secondly, for the smooth L1 loss, the gradient is dominated by the outliers with poor localization accuracy. In this work, IoU-balanced loss functions consisting of IoU-balanced classification loss and IoU-balanced localization loss are proposed to solve these problems. IoU-balanced classification loss pays more attention to positive examples with high IoU and enhances the correlation between classification and localization tasks. IoU-balanced localization loss decreases the gradient of examples with low IoU and increases the gradient of examples with high IoU, which improves the localization accuracy of models. Extensive experiments on MS COCO, PASCAL VOC, Cityscapes and WIDERFace demonstrate that IoU-balanced losses can substantially improve the popular single-stage detectors, especially the localization accuracy. On COCO test-dev, the proposed methods can substantially improve AP by 1.0%∼1.7% and AP75 by 1.0%∼2.4%. On PASCAL VOC, Cityscape and WIDERFace, it can also substantially improve AP by 1.0%∼1.5% and AP80, AP90 by ∼3.9%. The source code will be made publicly available.