NMS(解决大多数手写NMS运算效率较低的问题)

现在网上大多数自己手写的NMS都会导致GPU利用率较低，进而导致训练的时间较长，这是因为会经常调用单个box和其他box之间的IOU计算，导致内存和显存之间的IO交互较多。

def nms(boxes, scores, iou_threshold):
    keep = []
    idxs = scores.argsort()
    ious = circle_iou(boxes, boxes)
    while idxs.numel() > 0:
        max_score_index = idxs[-1]
        keep.append(max_score_index)
        if idxs.size(0) == 1:  # 就剩余一个框了；
            break
        idxs = idxs[:-1]
        idxs = idxs[ious[max_score_index, idxs] <= iou_threshold]
    keep = idxs.new(keep)
    return kee

标签：index,NMS,max,boxes,score,idxs,手写,运算
来源： https://www.cnblogs.com/yancx/p/16537680.html