facebookresearch / detectron2

Detectron2 is a platform for object detection, segmentation and other visual recognition tasks.
https://detectron2.readthedocs.io/en/latest/
Apache License 2.0
29.3k stars 7.32k forks source link

Focal Loss #5313

Open anishjv opened 1 week ago

anishjv commented 1 week ago

Hello FAIR team,

I am submitting a small pull-request, which adds focal loss (as detailed in https://arxiv.org/pdf/1708.02002v2) as a possible loss function in the Fast RCNN output layer (box predictor). The new commit simply adds an attribute to the class FastRCNNOutputLayers and makes the attribute accessible from within the LazyConfig workflow. Once a LazyConfig has been instantiated using cfg = LazyConfig.load('...') the attribute can be accessed by calling cfg.model.roi_heads.box_predictor.use_focal_ce = True. The attribute's value is 'False' by default.

Using focal loss for the ROI class prediction greatly improves network validation accuracy when training on datasets with a large deal of class imbalance. Datasets with large class imbalance are extremely common for researchers working with natural data such as time lapse movies of cells, where possible classes may be "mitotic" and "non-mitotic". The typical ratio of mitotic to non-mitotic cells on a frame is 1:9, and if training a RCNN to detect cells, whole frames must be used thus excluding the possibility of balancing the dataset by training on mitotic cells only for a select number of frames. For this reason loss functions such as focal loss are highly useful.