[LibTorch 입문] 7편: C++/Python 통합 모델 추론 파이프라인 실습

들어가며

이 시리즈에서 우리는 다음과 같은 단계를 거쳐왔습니다.

  • C++ 환경에서 LibTorch 사용법 익히기 (기초 텐서 연산, TorchScript 모델 로드)
  • Python에서 학습한 모델을 C++로 가져와 추론하기
  • pybind11을 통해 C++ 코드를 Python에 바인딩하기
  • C++과 Python 사이에서 텐서를 자유롭게 교환하는 기법 살펴보기

이제 여기까지 배운 내용을 종합하여, 하나의 일관된 파이프라인을 구축해봅시다. 최종적으로 다음과 같은 흐름을 구현할 예정입니다.

  1. Python에서 텐서(입력 데이터) 준비
  2. pybind11로 바인딩된 C++ 함수를 호출해 TorchScript 모델 추론 수행
  3. 결과 텐서를 Python으로 되돌려 받아 후처리 및 시각화

이 과정을 통해 C++ 성능과 Python의 편리함을 동시에 누리는 실제 파이프라인을 경험할 수 있습니다.

예제 시나리오

예를 들어, 다음과 같은 시나리오를 생각해봅시다.

  • 우리는 Python에서 이미지나 센서 데이터를 읽고, 전처리하여 (1, 3, 224, 224) 크기의 텐서로 만든다고 가정합니다.
  • 이 텐서를 C++에 전달해 TorchScript 모델(예: 분류 모델)을 사용해 추론하고, (1, 1000) 크기의 클래스 로짓(logits)을 반환받습니다.
  • Python에서는 이 로짓을 소프트맥스 및 argmax로 후처리하여 최종 클래스를 결정하고, 결과를 출력하거나 시각화합니다.

프로젝트 구조 예시

아래와 같은 디렉토리 구조를 가정하겠습니다.

my_project/
  ├─ CMakeLists.txt
  ├─ model.pt            # TorchScript로 변환한 모델 파일
  ├─ src/
  │   ├─ model_wrapper.cpp
  │   ├─ model_wrapper.h
  │   ├─ bind_module.cpp
  ├─ test_pipeline.py
  └─ README.md
  • model_wrapper.cpp/h: C++에서 TorchScript 모델 로딩 및 추론 로직을 구현한 클래스 제공
  • bind_module.cpp: pybind11 모듈로 ModelWrapper를 Python에 노출
  • test_pipeline.py: Python에서 전처리 → C++ 추론 호출 → 후처리 전체 흐름 테스트

C++ 코드 상세

ModelWrapper 클래스 (model_wrapper.h)

// model_wrapper.h
#pragma once
#include <torch/torch.h>
#include <torch/script.h>
#include <string>

class ModelWrapper {
public:
    ModelWrapper(const std::string& model_path);
    torch::Tensor infer(const torch::Tensor& input);
private:
    torch::jit::script::Module module_;
};

ModelWrapper 구현 (model_wrapper.cpp)

// model_wrapper.cpp
#include "model_wrapper.h"
#include <iostream>

ModelWrapper::ModelWrapper(const std::string& model_path) {
    try {
        module_ = torch::jit::load(model_path);
        std::cout << "Model loaded from " << model_path << "\n";
    } catch (const c10::Error& e) {
        std::cerr << "Error loading the model: " << e.what() << "\n";
        throw e;
    }
}

torch::Tensor ModelWrapper::infer(const torch::Tensor& input) {
    // 필요하면 GPU 사용
    if (torch::cuda::is_available()) {
        module_.to(torch::kCUDA);
    }

    auto in = input;
    if (torch::cuda::is_available()) {
        in = in.to(torch::kCUDA);
    }

    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(in);

    auto output = module_.forward(inputs).toTensor();

    if (output.device().is_cuda()) {
        output = output.cpu();
    }

    return output;
}

여기서 중요한 점:

  • 모델 파일 model.pt를 로딩하여 module_에 저장
  • infer 함수에서 입력 텐서를 받아 GPU가 가능하면 GPU로 보내고 추론 수행
  • 결과를 CPU로 다시 돌려 Python에 반환 (Python에서 후처리 편의를 위해 CPU 텐서로 바꿔줌)

pybind11 바인딩 (bind_module.cpp)

// bind_module.cpp
#include <pybind11/pybind11.h>
#include <torch/torch.h>
#include "model_wrapper.h"

namespace py = pybind11;

PYBIND11_MODULE(model_module, m) {
    py::class_<ModelWrapper>(m, "ModelWrapper")
        .def(py::init<const std::string&>())
        .def("infer", &ModelWrapper::infer, "Infer on input tensor");
}

이로써 Python에서 model_module를 import한 뒤 ModelWrapper("model.pt")로 인스턴스를 만들고, infer 메서드를 호출할 수 있습니다.

CMakeLists.txt 예제

cmake_minimum_required(VERSION 3.10)
project(MyFullPipeline)

set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

find_package(Torch REQUIRED)
find_package(pybind11 REQUIRED)

pybind11_add_module(model_module src/bind_module.cpp src/model_wrapper.cpp)
target_link_libraries(model_module "${TORCH_LIBRARIES}")
set_property(TARGET model_module PROPERTY CXX_STANDARD 11)

빌드 시:

cd my_project
mkdir build && cd build
cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
make

빌드가 완료되면 model_module.so(또는 .pyd) 파일이 생성됩니다.

Python 파이프라인 테스트 (test_pipeline.py)

import torch
import model_module

# 1. 전처리 (Python)
# 여기서는 가짜 데이터로 대체
# 실제 시나리오에서는 이미지 로드 & 정규화 등 수행
input_tensor = torch.randn(1, 3, 224, 224)

wrapper = model_module.ModelWrapper("model.pt")

# 2. C++ 추론 호출
output = wrapper.infer(input_tensor)

# 3. 후처리 (Python)
# 가정: output shape: (1, 1000) - 이미지 분류 결과
prob = torch.softmax(output, dim=1)
pred_class = torch.argmax(prob, dim=1)
print("Predicted class:", pred_class.item())

# 추가로 prob를 시각화하거나, 특정 클래스 레이블 맵핑 가능

여기서 input_tensor는 Python에서 준비한 텐서이고, wrapper.infer() 호출 시 C++의 TorchScript 모델 추론 로직이 수행됩니다. 반환된 output 텐서를 Python에서 받아 소프트맥스, argmax 등 후처리를 진행했습니다.

GPU 활용 확인

infer 함수 내에서 if (torch::cuda::is_available())를 통해 CUDA를 활용하도록 했으므로, GPU 환경이라면 자동으로 GPU 추론이 이뤄집니다. Python에서 input_tensor를 .cuda()로 올려 전달하면, C++도 GPU 텐서로 받아 GPU 추론 후 CPU로 결과를 돌려줍니다.

if torch.cuda.is_available():
    input_tensor = input_tensor.cuda()

output = wrapper.infer(input_tensor)

이렇게 하면 실제 GPU에서 추론이 돌아가며, Python 쪽에서 GPU-CPU 사이 텐서 복사 없이 매끄럽게 전송됩니다(단, infer 내에서 CPU 반환을 명시했으므로 최종 결과는 CPU 텐서로 돌아옴).

확장 아이디어

  • 다양한 입력 형식 지원:
    Python에서 이미지, 텍스트, 오디오 등을 텐서로 변환한 뒤 C++ 추론 호출 가능.
  • 배치 처리:
    한 번에 여러 입력을 담은 배치 텐서를 전달하고 C++에서 처리한 뒤 결과를 반환.
  • 비동기 처리:
    pybind11, std::future, async 등을 활용해 Python 스레드에서 비동기적으로 C++ 추론 수행.
  • 대규모 파이프라인 구축:
    이 원리를 확장해 Python 스크립트에서 데이터 로딩, 전처리, C++ 추론, 후처리, 모델 앙상블 등을 하나의 파이프라인으로 구성할 수 있음.

정리

이번 글에서는 C++과 Python 환경을 모두 활용하는 실제 추론 파이프라인 구축 예제를 살펴보았습니다. 여기까지 오면 다음과 같은 능력을 갖추게 됩니다.

  • Python의 편리함을 활용해 데이터 전처리, 후처리, 시각화 수행
  • C++의 성능과 LibTorch 모델 로딩/추론 기능을 활용해 고속 추론
  • pybind11을 통해 C++ 함수를 Python에서 자연스럽게 호출

이러한 구조는 연구 단계에서 Python으로 빠르게 실험한 뒤, 프로덕션 단계에서 C++ 기반으로 일부 로직을 고성능화하고, 여전히 Python 환경에서 컨트롤을 유지할 수 있게 합니다.

다음(마지막) 글에서는 전체 시리즈를 정리하고, 추가로 나아갈 수 있는 고급 주제나 최적화 전략에 대해 소개하며 마무리하겠습니다.

참고 자료

반응형