[LibTorch 입문] 6편: C++과 Python 사이에서 텐서 교환하기

 

앞선 글에서 우리는 C++에서 PyTorch 텐서를 생성하고 연산하며, TorchScript 모델을 C++에서 로드하는 방법, 그리고 pybind11을 통해 C++ 함수를 Python에서 호출하는 방법까지 살펴보았습니다. 이제 한 단계 더 나아가 C++과 Python 사이에서 텐서를 자유롭게 주고받는 방법을 다뤄보겠습니다.

이 과정은 다양한 시나리오에서 유용합니다. 예를 들어,

• Python에서 전처리한 입력 데이터를 C++ 모델 로직에 전달하고 싶을 때
• C++에서 계산한 텐서를 Python에서 시각화하거나 후처리하고 싶을 때
• Python에서 추론 로직 일부를 C++로 구현하여 성능을 향상시키고, 그 결과를 다시 Python으로 반환할 때

이 글에서는 pybind11과 LibTorch를 활용해 C++과 Python 텐서를 교환하는 기법을 예제 코드와 함께 안내하겠습니다.

C++과 Python 텐서의 기본 호환성

PyTorch Python API에서 사용하는 torch.Tensor는 내부적으로 C++의 torch::Tensor 객체를 래핑(wrap)한 것입니다. 이를 기반으로 PyTorch는 pybind11 매커니즘을 통해 C++ 텐서와 Python 텐서를 자연스럽게 변환합니다. 즉, pybind11로 만든 C++ 함수에 torch::Tensor를 인자로 받거나 반환하도록 하면, Python에서 해당 함수를 torch.Tensor로 호출하고 결과를 torch.Tensor로 받을 수 있습니다.

이 덕분에 별도의 변환 코드를 작성할 필요 없이 C++과 Python 사이에서 텐서를 교환할 수 있습니다. 앞서 5편에서 이미 tensor_add 예제로 이런 교환이 가능함을 확인했습니다.

기본적인 텐서 교환 예제

이미 앞 글에서 유사한 예제를 다뤘지만, 여기서 다시 간단한 예제로 시작해 보겠습니다. C++에서 두 텐서를 받아서 더한 뒤 반환하는 함수를 Python에서 호출하는 예제입니다.

// tensor_ops.cpp
#include <torch/torch.h>

// 두 텐서를 element-wise로 더하는 함수
torch::Tensor add_tensors(const torch::Tensor& a, const torch::Tensor& b) {
    return a + b;
}

// bind_ops.cpp
#include <pybind11/pybind11.h>
#include <torch/torch.h>

namespace py = pybind11;

extern torch::Tensor add_tensors(const torch::Tensor& a, const torch::Tensor& b);

PYBIND11_MODULE(tensor_ops, m) {
    m.doc() = "Tensor operation module";
    m.def("add_tensors", &add_tensors, "Add two tensors");
}

빌드 후 Python에서 테스트:

import torch
import tensor_ops

a = torch.randn(2, 2)
b = torch.randn(2, 2)
c = tensor_ops.add_tensors(a, b)
print("a+b:", c)

이처럼 별도의 변환 없이 Python 텐서를 C++ 함수에 전달하고, C++ 함수가 반환한 텐서를 Python에서 곧바로 사용할 수 있습니다.

텐서 장치(Device) 처리

C++/Python 간 텐서 교환 시 주의할 점 중 하나는 텐서가 CPU인지 GPU(CUDA)인지에 대한 장치(device) 정보입니다. Python에서 CUDA 텐서를 C++ 함수에 전달하면, C++에서도 해당 텐서를 GPU 메모리에 올린 상태로 받습니다. 반대로 C++에서 GPU 텐서를 반환하면 Python에서도 그대로 GPU 텐서를 받게 됩니다.

이러한 기작 덕분에 장치를 일일이 강제 변환할 필요가 없습니다. 다만, 환경에 따라 CUDA가 없는 상황에서 CUDA 텐서를 반환하는 경우 에러가 발생할 수 있으므로, C++ 측 코드에서 torch::cuda::is_available() 등을 사용해 GPU 장치 지원 여부를 확인하고 적절히 핸들링하는 것이 좋습니다.

예제: C++에서 텐서를 GPU로 옮긴 뒤 반환하기

torch::Tensor process_on_gpu(const torch::Tensor& input) {
    if (torch::cuda::is_available()) {
        auto gpu_input = input.to(torch::kCUDA);
        // 어떤 연산 수행
        auto result = gpu_input * 2;
        return result; // GPU 텐서 반환
    } else {
        // fallback: CPU에서 처리
        return input * 2;
    }
}

Python 측에서:

import torch
import tensor_ops

x = torch.randn(2, 2)
if torch.cuda.is_available():
    x = x.cuda()

y = tensor_ops.process_on_gpu(x)
print(y.device)  # GPU 텐서이면 cuda:0 출력
print(y)

텐서 메모리 공유

Python과 C++ 사이에서 텐서를 교환할 때, 텐서의 메모리가 복사 없이 공유될 수 있다는 점도 중요한 최적화 포인트입니다. PyTorch 텐서는 참조 카운팅(reference counting)을 통해 메모리 관리하기 때문에, C++과 Python 양쪽에서 동일한 텐서 메모리를 공유할 수 있습니다.

즉, C++에서 받은 텐서의 내용을 Python에서 수정하면 C++ 쪽에도 반영될 수 있고, 그 반대도 가능합니다. 다만, 이러한 동작은 로우레벨 메모리 관리와 관련된 이슈를 초래할 수 있으므로, 가능한 한 함수형 스타일(입력 텐서는 변경하지 않고 새 텐서 반환)로 처리하는 것을 권장합니다.

고유한 텐서 변환 로직이 필요한 경우

일반적으로 torch::Tensor는 자동 변환이 가능하지만, 만약 특정한 C++ 텐서 타입을 Python으로 변환하고 싶다면 별도의 변환 로직이 필요할 수 있습니다. 예를 들어, Eigen 텐서나 OpenCV 이미지 등을 torch::Tensor로 바꾼 뒤 Python에 전달하거나, Python에서 받아온 torch.Tensor를 Eigen 행렬로 변환하는 등의 작업입니다.

이를 위해서는 다음과 같은 단계를 거칠 수 있습니다.

1. Python에서 torch.Tensor로 데이터를 받음
2. C++ 함수 내에서 tensor.data_ptr<type>() 등을 이용해 원시 포인터 접근
3. 원하는 형식(Eigen 행렬 등)으로 복사 또는 map
4. 처리 후 다시 torch::Tensor로 변환해 반환

예를 들어, C++에서 Eigen을 사용한다면 다음과 같이 할 수 있습니다.

#include <torch/torch.h>
#include <Eigen/Dense>

Eigen::MatrixXd tensor_to_eigen(const torch::Tensor& t) {
    // 예: double형 텐서라고 가정
    TORCH_CHECK(t.dtype() == torch::kFloat64, "Expected float64 tensor");
    TORCH_CHECK(t.is_contiguous(), "Expected contiguous tensor");

    auto sizes = t.sizes();
    // 2차원 텐서라고 가정
    TORCH_CHECK(sizes.size() == 2, "Expected 2D tensor");
    int rows = sizes[0];
    int cols = sizes[1];

    const double* data = t.data_ptr<double>();
    Eigen::Map<const Eigen::Matrix<double, Eigen::Dynamic, Eigen::Dynamic, Eigen::RowMajor>> mat(data, rows, cols);
    return Eigen::MatrixXd(mat); // 복사
}

torch::Tensor eigen_to_tensor(const Eigen::MatrixXd& mat) {
    auto options = torch::TensorOptions().dtype(torch::kFloat64);
    torch::Tensor t = torch::from_blob((void*)mat.data(), {mat.rows(), mat.cols()}, options).clone();
    return t; // 복사된 텐서 반환
}

이런 식으로 중간 변환 과정을 거쳐 C++과 Python을 오가며 다양한 타입의 데이터를 다룰 수 있습니다.

실전 응용 예제: Python 전처리 → C++ 추론 → Python 후처리

1. Python에서 데이터 로드 및 전처리: Python으로 이미지를 로드하고, torch.Tensor 형태로 전처리를 수행한 뒤 C++ 함수에 전달.
2. C++에서 TorchScript 모델 추론: Python에서 받은 텐서를 C++ 함수로 넘겨 TorchScript 모델을 실행, 결과 텐서를 반환.
3. Python 후처리 및 시각화: C++에서 반환한 텐서를 Python에서 받아 후처리(softmax나 argmax, 시각화 등)를 수행.

예제 스니펫:

# Python code
import torch
import model_module  # C++ 바인딩된 모듈, ModelWrapper 예제
import matplotlib.pyplot as plt

# 이미지 로드 & 전처리 (Python)
img_tensor = torch.randn(1, 3, 224, 224)  # 예시: 임의 입력

wrapper = model_module.ModelWrapper("model.pt")
output = wrapper.infer(img_tensor)  # C++에서 추론 수행

# Python 후처리: 소프트맥스, argmax 등
prob = torch.softmax(output, dim=1)
pred_class = torch.argmax(prob, dim=1)
print("Predicted class:", pred_class.item())

여기서 infer 함수는 C++에서 TorchScript 모델을 로드해 추론하는 로직을 담고 있으며, Python에서 넘긴 텐서를 그대로 받아 추론 후 반환합니다. 모든 텐서 교환은 pybind11을 통한 자동 변환 덕분에 번거로운 형 변환 없이 이뤄집니다.

정리

이번 글에서는 C++과 Python 사이에서 텐서를 교환하는 방법과 주의사항, 그리고 실전 응용 예제를 살펴보았습니다.

• 기본 호환성: torch::Tensor는 pybind11을 통해 Python과 C++ 사이에서 자연스럽게 교환 가능
• 장치 관리: CPU/GPU 텐서 모두 자동 변환 가능, 필요 시 device 체크
• 메모리 공유와 커스텀 변환: 텐서 메모리 공유로 복사 비용 최소화, 필요하면 커스텀 변환 로직 작성 가능
• 실전 예제: Python 전처리 → C++ 추론 → Python 후처리 파이프라인 구축

다음 글에서는 이러한 텐서 교환과 모델 추론 로직을 종합하여, C++과 Python 환경을 유기적으로 통합한 실제 파이프라인을 구성하는 더 큰 예제를 다루며, 시리즈를 마무리 단계로 이끌어가겠습니다.

참고 자료

• PyTorch C++ Docs
• pybind11 공식 문서
• PyTorch to C++ Tensor API