nms.cpp 코드 분석

• 해당 파일은 Python에서 넘어온 tensor를 확인하고, 점수순으로 정렬한 뒤 CUDA 함수로 넘김

---

include 부분

#include <torch/extension.h>
PyTorch C++ extension을 만들기 위한 헤더


#include <torch/types.h>
at::Tensor 같은 PyTorch tensor 타입을 쓰기 위한 헤더


#include <iostream>
C++ 입출력용 헤더

---

CUDA 함수 선언

std::vector<at::Tensor> nms_cuda_forward(
        at::Tensor boxes,
        at::Tensor idx,
        float nms_overlap_thresh,
        unsigned long top_k);

• 이 함수는 nms_kernel.cu에 실제 구현이 있음
• 하지만 nms.cpp에서 이 함수를 호출하려면 먼저 이런 함수가 있다고 알려줘야 함 따라서 함수 선언을 해둠
• 반환값은 std::vector<at::Tensor>
  • PyTorch tensor 여러 개를 vector로 묶어서 반환
• 인자는 다음과 같음

boxes
차선 후보 tensor

idx
점수 높은 순서로 정렬된 index tensor

nms_overlap_thresh
두 차선이 비슷한지 판단하는 기준값

top_k
최대 몇 개 차선을 남길지

---

입력 체크 매크로

CHECK_CUDA

#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")

• x가 GPU tensor인지 확인함
• #define은 C/C++의 매크로, 맨뒤에 ;를 안씀
• AT_ASSERTM은 PyTorch C++/CUDA extension 코드에서 쓰는 조건 검사 매크로

AT_ASSERTM(조건, "에러 메시지")

• 조건이 참이면 그대로 진행
• 조건이 거짓이면 프로그램을 멈추고 에러 메시지를 출력

x.type().is_cuda()

• x라는 PyTorch tensor의 타입 정보를 가져옴
• 이후 tensor가 CUDA tensor인지 확인함

---

CHECK_CONTIGUOUS

#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")

• tensor가 메모리상에서 연속적으로 저장되어 있는지 확인함
• .is_contiguous() PyTorch Tensor가 메모리에 연속적으로 저장되어 있는지 확인하는 함수

---

CHECK_INPUT

#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)

• CUDA tensor인지 확인하고, contiguous인지도 확인하는 코드
• 위에서 정의한 CHECK_CUDA(x), CHECK_CONTIGUOUS(x)를 둘 다 실행함

---

nms_forward 함수

정의 부분

std::vector<at::Tensor> nms_forward(
        at::Tensor boxes,
        at::Tensor scores,
        float thresh,
        unsigned long top_k) {

• 인자 의미는 다음과 같다.

boxes
차선 후보들
shape: [N, 77]

scores
각 차선 후보의 점수
shape: [N]

thresh
NMS threshold

top_k
최대 keep 개수

---

점수 정렬

auto idx = std::get<1>(scores.sort(0,true));

• scores를 내림차순으로 정렬하고, 정렬된 값 말고 정렬된 index만 꺼내서, idx라는 변수에 저장함
• 아래에 코드 하나씩 분해함

scores.sort(0, true)

예시
scores = [0.2, 0.9, 0.7]

scores.sort(0, true)
= (
    values  = [0.9, 0.7, 0.2],
    indices = [1, 2, 0]
  )

• PyTorch C++ Tensor API
• scores는 C++의 at::Tensor
• 0번 차원 기준으로 내림차순으로 정렬(False면 오름차순)
• PyTorch의 sort는 결과를 두 개 반환 => 정렬된 점수 values, 정렬된 원래 위치 indices

std::get<1>(...)

• 표준 C++ 기능
• C++의 std::tuple에서 몇 번째 값을 꺼낼 때 사용
• scores.sort()가 반환한 두 값 중 두 번째 값을 꺼내라는 뜻

---

입력 체크

CHECK_INPUT(boxes);
CHECK_INPUT(idx);

• boxes와 idx가 CUDA tensor이고 contiguous인지 확인함
• 이때 scores 자체는 여기서 체크하지 않음

---

CUDA 함수 호출

return nms_cuda_forward(boxes, idx, thresh, top_k);

• 정렬된 idx를 포함하여 NMS에 필요한 정보를 실제 CUDA 쪽 함수로 넘김
• nms_kernel.cu 에서 처리할 예정

---

PYBIND11_MODULE

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("nms_forward", &nms_forward, "NMS");
}

• C++ 함수를 Python에서 호출할 수 있게 등록하는 코드
• C++ 함수 nms_forward를 Python 모듈 안의 nms_forward라는 이름으로 등록

PYBIND11_MODULE
C++/CUDA 코드를 Python 모듈로 노출하는 pybind11 매크로

TORCH_EXTENSION_NAME
setup.py의 CUDAExtension name에서 온 모듈 이름
이 프로젝트에서는 nms_impl

m
Python 모듈 객체

m.def(...)
Python에서 호출 가능한 함수를 등록

"nms_forward"
Python에서 보이는 함수 이름

&nms_forward
연결할 C++ 함수 주소

"NMS"
함수 설명 문자열