[LibTorch 입문] 4편: Python 모델을 C++에서 TorchScript로 추론하기

 

많은 개발자와 연구자들이 PyTorch를 이용해 Python 환경에서 딥러닝 모델을 개발하고 학습합니다. 하지만 실제 프로덕션 환경이나 고성능 애플리케이션에서는 C++ 기반의 애플리케이션에 모델을 통합하고 싶을 때가 있습니다. 이때 Python 환경 없이도 모델을 로딩하고 추론할 수 있도록 해주는 것이 바로 TorchScript 입니다.

TorchScript를 사용하면 Python으로 학습한 PyTorch 모델을 별도의 .pt 파일 형태로 내보내고, 이 파일을 C++ LibTorch 환경에서 로딩해 추론할 수 있습니다. 이 글에서는 Python에서 TorchScript 모델을 만드는 방법, 그리고 C++에서 이를 로딩해 추론하는 과정을 단계별로 살펴봅니다. 또한 단순한 완전연결 모델에서 한 걸음 더 나아가, 입력 형태가 다른 예제도 해보고, GPU 활용과 예외 처리 등의 주제도 다룹니다.

TorchScript란 무엇이며 왜 필요할까?

일반적으로 PyTorch 모델은 Python 인터프리터에서 동작하며, Python 코드를 통해 모델 구조를 정의하고 학습합니다. 이러한 Python 종속성을 제거하고, 모델을 독립적인 형태로 만들기 위해 PyTorch는 TorchScript라는 중간 표현(intermediate representation, IR)을 지원합니다. TorchScript 모델은 다음과 같은 이점을 제공합니다.

• Python 환경 불필요: .pt 파일 형태로 모델을 저장하면 Python 없이 C++ LibTorch 만으로 모델 로딩 및 추론이 가능합니다.
• 성능 최적화 및 이동성: 모델을 C++ 애플리케이션에 직접 통합함으로써 언어 간 오버헤드를 줄이고, 성능을 극대화할 수 있습니다.
• 다양한 플랫폼 지원: Python 런타임이 없는 환경(예: 임베디드 시스템, 컨테이너)에서도 모델을 사용할 수 있습니다.

TorchScript 모델을 만들기 위해서는 크게 두 가지 방법이 있습니다.

1. torch.jit.trace(): 입력 텐서를 모델에 통과시키며 연산 그래프를 "추적"하여 TorchScript 모델을 생성하는 방식. 주로 제어 흐름이 고정된 모델에 적합합니다.
2. torch.jit.script(): 모델의 Python 코드 자체를 스크립팅하여 TorchScript로 변환. 조건문, 반복문 등 동적 제어 흐름도 처리 가능.

입문 단계에서는 trace 방식이 간단하므로 이를 사용합니다.

Python에서 모델 내보내기: torch.jit.trace 사용 예제

다음은 Python에서 간단한 완전연결 신경망(Linear Layer) 모델을 정의하고, trace를 이용해 .pt 모델 파일로 내보내는 예시입니다. 이 모델은 입력차원 4, 출력차원 2를 가지는 단순한 구조입니다.

# save_model.py
import torch
import torch.nn as nn

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(4, 2)  # 입력: 4차원, 출력: 2차원

    def forward(self, x):
        return self.fc(x)

if __name__ == "__main__":
    model = SimpleModel()
    model.eval()  # 추론 모드로 전환 (BatchNorm, Dropout 고정화)
    
    # 더미 입력 텐서 (1x4)
    example_input = torch.randn(1, 4)

    # trace를 통해 모델 그래프 추출
    traced_model = torch.jit.trace(model, example_input)
    
    # TorchScript 모델 파일로 저장
    traced_model.save("model.pt")
    print("model.pt saved successfully!")

이 코드를 실행하면 현재 디렉토리에 model.pt 파일이 생성됩니다. 여기에는 모델 구조와 가중치 정보가 포함되어 있으며, C++에서 로딩 가능하도록 정리되어 있습니다.

더 복잡한 모델 예시

만약 CNN(합성곱 신경망)이나 RNN 등을 사용하고 싶다면 유사한 방식으로 trace할 수 있습니다. 단, trace 방식은 주어진 입력 텐서로 실행되는 정적 경로만 기록하기 때문에, 입력 형태나 크기가 변하는 모델이라면 script 방식을 고려해야 합니다.

예를 들어, 간단한 CNN 모델을 TorchScript로 저장하고 싶다면 다음과 같이 할 수 있습니다.

import torch
import torch.nn as nn

class SimpleCNN(nn.Module):
    def __init__(self):
        super(SimpleCNN, self).__init__()
        self.conv = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)
        self.fc = nn.Linear(16*6*6, 10)  # 예: 입력 이미지가 3x8x8이라 가정했을 때

    def forward(self, x):
        x = self.conv(x)
        x = torch.relu(x)
        x = torch.flatten(x, 1)
        x = self.fc(x)
        return x

if __name__ == "__main__":
    model = SimpleCNN().eval()
    example_input = torch.randn(1, 3, 8, 8)
    traced_model = torch.jit.trace(model, example_input)
    traced_model.save("cnn_model.pt")

이렇게 하면 3x8x8 이미지를 입력받아 10차원 출력으로 매핑하는 CNN 모델도 TorchScript 형태로 저장할 수 있습니다.

C++에서 모델 로드하기

이제 C++ 코드에서 model.pt를 로드하고 추론해봅시다. 기본 구조는 다음과 같습니다.

1. #include <torch/script.h>로 TorchScript 관련 헤더를 포함
2. torch::jit::load("model.pt")를 통해 모델 로딩
3. 입력 텐서를 준비하고, module.forward({inputs}) 호출
4. 결과 텐서 출력

기본 예제 코드

#include <torch/torch.h>
#include <torch/script.h>
#include <iostream>

int main() {
    try {
        // model.pt 로드
        torch::jit::script::Module module = torch::jit::load("model.pt");
        std::cout << "Model loaded successfully.\n\n";

        // 추론에 사용할 입력 텐서 (1x4)
        torch::Tensor input = torch::randn({1, 4});
        std::cout << "Input:\n" << input << "\n\n";

        // forward 실행
        std::vector<torch::jit::IValue> inputs;
        inputs.push_back(input);
        torch::Tensor output = module.forward(inputs).toTensor();
        std::cout << "Output:\n" << output << "\n\n";
    }
    catch (const c10::Error& e) {
        std::cerr << "Error loading or running the model: " << e.what() << "\n";
        return -1;
    }

    return 0;
}

위 코드 실행 시, 랜덤 입력 텐서에 대한 모델의 출력 텐서가 콘솔에 찍히게 됩니다. 여기서 모델의 가중치는 save_model.py 실행 시점의 상태입니다. 실제 학습된 가중치를 가진 모델이라면 의미있는 추론 결과를 얻을 수 있을 것입니다.

예외 처리 및 경로 지정

• 모델 경로가 잘못되었을 경우 c10::Error 예외가 발생할 수 있습니다.
• GPU 환경에서 CUDA를 활용할 경우, 모델과 텐서를 GPU로 .to(torch::kCUDA)를 통해 옮기고 추론할 수 있습니다.

예를 들어:

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

이렇게 하면 GPU에서 추론을 수행하게 되어 대형 모델일 경우 속도 향상에 유리합니다.

다양한 입력 형태 처리하기

trace 방식은 모델을 변환할 때 사용한 입력 텐서 형태에 최적화됩니다. 즉, 추론 시에도 동일하거나 호환 가능한 입력 크기를 사용해야 합니다. 만약 추후 입력 크기가 달라질 수 있다면, torch.jit.script()를 사용하거나, nn.AdaptiveAvgPool2d와 같은 크기 불변 연산을 활용하여 입력 형상의 유연성을 확보할 수 있습니다.

# 예: script를 사용해 변환
scripted_model = torch.jit.script(model)
scripted_model.save("model_scripted.pt")

C++에서는 model_scripted.pt를 동일하게 load하여 사용하면 됩니다.

추가적인 코드 예제: CNN 모델 추론

아래는 앞서 Python에서 만들었던 cnn_model.pt를 C++에서 로딩하는 예시입니다.

#include <torch/script.h>
#include <iostream>

int main() {
    try {
        torch::jit::script::Module module = torch::jit::load("cnn_model.pt");
        std::cout << "CNN Model loaded successfully.\n";

        // 1x3x8x8 이미지 텐서
        auto input = torch::randn({1, 3, 8, 8});
        
        // GPU 사용 가능 시 GPU 이동
        if (torch::cuda::is_available()) {
            std::cout << "CUDA available! Moving model and input to GPU.\n";
            module.to(torch::kCUDA);
            input = input.to(torch::kCUDA);
        }

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

        torch::Tensor output = module.forward(inputs).toTensor();

        // 출력 차원: [1, 10] 의 로짓(logit)
        std::cout << "CNN Output:\n" << output << "\n";
    }
    catch (const c10::Error& e) {
        std::cerr << "Error: " << e.what() << "\n";
        return -1;
    }

    return 0;
}

이 코드에서는 1x3x8x8 텐서를 입력으로 넣어 10차원 출력 벡터를 얻습니다. GPU가 가능하면 모델과 입력을 GPU로 옮겨 더 빠른 추론이 가능합니다.

정리

지금까지 Python에서 학습한 PyTorch 모델을 TorchScript를 통해 .pt로 내보내고, 이를 C++ LibTorch 환경에서 로딩하여 추론하는 전체 과정을 살펴보았습니다.

• Python 측면: 모델 정의 → model.eval() 모드 전환 → torch.jit.trace() 또는 torch.jit.script() → .pt 파일로 저장
• C++ 측면: torch::jit::load()로 모델 로딩 → 입력 텐서 준비 → module.forward(inputs) 호출 → 출력 텐서 확인

이 과정을 통해 C++ 환경에서 딥러닝 모델을 사용할 수 있으며, Python 종속성이 사라져 다양한 배포 및 최적화 전략을 시도할 수 있습니다. 또한 GPU 활용, 다양한 모델 구조 지원, 동적 입력 크기 처리 등 확장 가능한 방법들이 많습니다.

향후에는 pybind11을 통해 C++ 로직을 Python에서 다시 호출 가능하도록 연결하거나, 더 복잡한 모델 및 최적화 전략을 적용하는 등의 추가 단계로 나아갈 수 있습니다.

참고 자료

• PyTorch 공식 홈페이지
• TorchScript 튜토리얼
• C++ TorchScript API 문서
• YouTube: PyTorch Tutorial