TorchScript와 C++에서 커스텀 연산 지원하기

📝 들어가며

PyTorch에서 Python으로 학습한 모델을 TorchScript로 변환하면, C++(LibTorch) 환경에서도 모델을 로드해 추론할 수 있습니다.
하지만 모델 내부에 torch_scatter, torch_fpsample 같은 커스텀 연산(custom operator) 이 포함되어 있다면, 단순히 .pt 파일만 로드하는 것으로는 부족합니다.

이 글에서는 torch_scatter, torch_fpsample 사례를 통해 TorchScript와 C++에서의 커스텀 연산 지원 조건과 방법을 정리했습니다.

---

🔹 TorchScript와 torch_scatter

• torch_scatter는 PyTorch 확장 라이브러리로, 이미 C++/CUDA 연산 등록이 되어 있습니다.
• 따라서 TorchScript로 변환된 모델 안에서도 해당 연산 노드가 유지되고, C++ 환경에 torch_scatter의 C++ 확장 라이브러리(libtorchscatter) 가 설치·로드되면 정상적으로 동작합니다.
• 만약 라이브러리를 로드하지 않으면 Unknown builtin op: torch_scatter::... 오류가 발생합니다.
• ✅ 즉, TorchScript + C++ 추론이 가능하지만 환경에 libtorchscatter.so/.dll을 반드시 배포해야 합니다.

---

🔹 TorchScript와 torch_fpsample

• torch_fpsample은 Farthest Point Sampling(FPS) 연산을 제공하는 확장 라이브러리입니다.
• 현재 구현 상태는 CPU 전용, 개발 진행 중으로, TorchScript/LibTorch에서의 C++ 연산 등록 지원이 명확하지 않습니다.
• 따라서 그대로 TorchScript 모델에 포함시켜 C++에서 바로 실행할 수 있다고 보장하기 어렵습니다.
• 대안으로는 PyTorch3D의 sample_farthest_points 같은 TorchScript 호환 FPS 연산을 사용하는 것이 안정적입니다.

---

🔹 TorchScript에서 커스텀 연산을 쓰려면?

1. C++ 연산 등록

   • TORCH_LIBRARY 매크로를 사용해 네임스페이스와 연산 시그니처를 등록해야 합니다.

   • 예시:

     TORCH_LIBRARY(my_ops, m) {
         m.def("warp_perspective(Tensor img, Tensor M) -> Tensor", warp_perspective);
     }

2. 공유 라이브러리 빌드

   • CMake로 .so/.dll 형태로 빌드하고, Python에서는 torch.ops.load_library로, C++에서는 런타임에 로드되도록 설정해야 합니다.

3. TorchScript 내보내기

   • Python에서 모델을 스크립팅/트레이싱할 때 해당 연산을 실제 호출해 .pt에 포함시킵니다.

4. C++ 실행

   • torch::jit::load("model.pt")로 모델을 불러오기 전, 커스텀 연산 라이브러리가 반드시 로드되어 있어야 합니다.

---

🔹 실무 체크리스트

• PyTorch와 LibTorch의 버전 및 ABI, CUDA 버전 일치 확인
• Python에서 torch.ops.load_library 호출 후 TorchScript 변환이 정상 동작하는지 검증
• C++ 실행 시 커스텀 연산 라이브러리(.so/.dll)가 로드되는지 확인
• 오류 발생 시 가장 먼저 “Unknown builtin op” → 연산 등록 문제부터 점검

---

✅ 결론

• torch_scatter: TorchScript 및 C++에서 지원 가능. 단, C++ 런타임 환경에 torch-scatter 확장 라이브러리가 설치·로드되어 있어야 함.
• torch_fpsample: TorchScript/LibTorch 지원이 아직 불확실. 안정적인 FPS 연산이 필요하다면 PyTorch3D 같은 대안 검토 권장.
• 일반 원칙: TorchScript에서 커스텀 연산을 활용하려면, 반드시 C++ 연산 등록 + 공유 라이브러리 로드가 되어 있어야 한다.

---