현재 Maskrcnn을 c++로 구현해보고 있습니다. 그 과정에서 nms를 TensorRT API를 사용하지 않고
cuda로 구현해보고 있는데 결과가 일치하지 않아 문의드립니다.
제가 구현한 nms 방식은 parallel nms로 모든 bbox를 돌며 iou를 계산 후 iou_thres 값보다 큰 경우 conf 값이 더 작은 bbox를 없애는 형식으로 구현하였습니다.
하지만 python에서 maskrcnn을 실행할 경우 제 로직에선 지워지는 bbox가 지워지지 않더라구요.
그래서 torchvision.ops.nms의 구체적인 로직이 궁금합니다.
도와주신다면 감사하겠습니다.
nms 로직을 제가 직접 구현해본 것은 아니라 도움이 안될것 같습니다.
대신 vision/nms_kernel.cpp at main · pytorch/vision · GitHub 링크에 가시면
torchvision.ops.nms 에 대해서 구현해 놓은 코드가 있으니 보시면서 본인 코드와 차이를 비교해 보셔야 할 것 같습니다.
torchvision.ops.nms는 Non-Maximum Suppression 알고리즘을 구현한 함수로, 주로 object detection에서 overlapping 된 bbox들 중 가장 정확도가 높은 bbox를 선택하는 용도로 사용됩니다. 이 함수 내부적으로 다양한 기술을 사용하여 NMS 알고리즘을 최적화하고, GPU로 실행될 수 있도록 텐서 형식 자료를 다룹니다.
보통 NMS 알고리즘은 3단계로 이루어집니다. 먼저, bounding box들을 confidence score에 따라 내림차순으로 정렬합니다. 이후, 가장 confidence score가 높은 bbox를 선택하고, 나머지 bbox들과 IoU (Intersection over Union) 값이 일정 이상인 bbox를 제거합니다. 마지막 단계로, 제거된 bbox들을 제외하고 남은 bbox들을 대상으로 똑같은 과정을 반복하여 결국 정해진 임계값 (IoU_threshold) 이상 overlap 되는 bbox 쌍들 중 더 작은 confidence score를 가진 bbox를 제거합니다.
따라서, torchvision.ops.nms 내부적으로 구현된 로직은 다음과 같습니다.
def nms(dets, scores, iou_threshold):
# dets: bounding box 좌표 (x1,y1,x2,y2)
# scores: bounding box에 대한 confidence score
# iou_threshold: IoU 임곗값
# 1. confidence score에 따라 bbox를 내림차순으로 정렬합니다.
_, order = scores.sort(0, descending=True)
# 2. bbox overlap을 계산
keep = []
while order.numel() > 0:
i = order[0]
keep.append(i)
if order.numel() == 1:
break
# overlapping 된 bbox들 중에서 가장 confidence score가 작은 bbox를 제거합니다.
# IoU 값이 일정 이상이면 제거합니다.
ovr = box_iou(dets[i], dets[order[1:]])
inds = torch.nonzero(ovr <= iou_threshold).squeeze()
if inds.numel() == 0:
break
# order, dets, scores에서 overlap 된 bbox 제거합니다.
order = order[inds + 1]
return torch.LongTensor(keep)
위의 코드에서는 먼저 주어진 bbox(dets)에 대한 confidence score(scores)에 따라 내림차순으로 정렬합니다. 그리고 나서, bbox들 간 overlapping을 계산하여 일정 이상 겹치는 경우를 제거하는 방식으로 NMS 알고리즘을 실행합니다. 결과적으로 제거된 bbox들이 계산되고, 남은 bbox들이 반환됩니다.
참고할만한 문서로는 다음의 논문과 코드가 있습니다: