torchvision.ops.nms 구체적인 로직이 궁금합니다

묻고 답하기
,
#1

현재 Maskrcnn을 c++로 구현해보고 있습니다. 그 과정에서 nms를 TensorRT API를 사용하지 않고
cuda로 구현해보고 있는데 결과가 일치하지 않아 문의드립니다.

제가 구현한 nms 방식은 parallel nms로 모든 bbox를 돌며 iou를 계산 후 iou_thres 값보다 큰 경우 conf 값이 더 작은 bbox를 없애는 형식으로 구현하였습니다.

하지만 python에서 maskrcnn을 실행할 경우 제 로직에선 지워지는 bbox가 지워지지 않더라구요.
그래서 torchvision.ops.nms의 구체적인 로직이 궁금합니다.

도와주신다면 감사하겠습니다.

#2

nms 로직을 제가 직접 구현해본 것은 아니라 도움이 안될것 같습니다.

대신 vision/nms_kernel.cpp at main · pytorch/vision · GitHub 링크에 가시면
torchvision.ops.nms 에 대해서 구현해 놓은 코드가 있으니 보시면서 본인 코드와 차이를 비교해 보셔야 할 것 같습니다.