🚀 Feature
The postprocessing step in RetinaNet is slow, and the whole inference time for RetinaNet is almost twice slower than Faster R-CNN as of today.
In particular,
|
for class_index in range(num_classes): |
|
# remove low scoring boxes |
|
inds = torch.gt(scores_per_image[:, class_index], self.score_thresh) |
|
boxes_per_class, scores_per_class, labels_per_class = \ |
|
boxes_per_image[inds], scores_per_image[inds, class_index], labels_per_image[inds, class_index] |
|
other_outputs_per_class = [(k, v[inds]) for k, v in other_outputs_per_image] |
|
|
|
# remove empty boxes |
|
keep = box_ops.remove_small_boxes(boxes_per_class, min_size=1e-2) |
|
boxes_per_class, scores_per_class, labels_per_class = \ |
|
boxes_per_class[keep], scores_per_class[keep], labels_per_class[keep] |
|
other_outputs_per_class = [(k, v[keep]) for k, v in other_outputs_per_class] |
|
|
|
# non-maximum suppression, independently done per class |
|
keep = box_ops.nms(boxes_per_class, scores_per_class, self.nms_thresh) |
|
|
|
# keep only topk scoring predictions |
|
keep = keep[:self.detections_per_img] |
|
boxes_per_class, scores_per_class, labels_per_class = \ |
|
boxes_per_class[keep], scores_per_class[keep], labels_per_class[keep] |
|
other_outputs_per_class = [(k, v[keep]) for k, v in other_outputs_per_class] |
|
|
|
image_boxes.append(boxes_per_class) |
|
image_scores.append(scores_per_class) |
|
image_labels.append(labels_per_class) |
|
|
|
for k, v in other_outputs_per_class: |
|
if k not in image_other_outputs: |
|
image_other_outputs[k] = [] |
|
image_other_outputs[k].append(v) |
does a for loop over the number of classes. This for loop can be parallelized by batching operations together over all classes, which should greatly improve the inference speed.
For reference, Detectron2 has sped up inference on RetinaNet a few times already, with latest optimization present in facebookresearch/detectron2@8999946 , and also batch inference over classes (and only does a for loop on the number of feature maps, which is much smaller than the number of COCO classes)
