파이토치 모델을 C++에서 사용하기 -1-

공식홈페이지 튜토리얼 정리글입니다

https://pytorch.org/tutorials/advanced/cpp_export.html

STEP 1: 파이토치 모델을 Torch script로 컨버팅

 Torch Script 을 이용해서 파이토치모델을 C++에서 불러올 수 있는데 여기엔 두가지 방법이 있습니다.

첫번째는 tracing 이라고 불리는 것으로 예제 인풋을 이용해서 그 인풋이 모델을 통과하면서 변화과정을 기록하는 것으로 사용 방법(control flow)이 제한적인 모델들에 적합한 방법

두번째는 모델에 소위 '주석'을 달아 Torch scrip 컴파일러가 직접적으로 파싱하고 컴파일할 수 있도록 하는 방법으로 위 Torch Script 링크를 누르면 자세한 방법이 나옴

Converting to Torch Script via Tracing

torch.jit.trace 함수에 모델 인스턴스와 예제 인풋값을 전달해주면 됨. 이러면 torch.jit.ScriptModule 오브젝트가 생성되며 모듈의 forward 메소드를 추적함

import torch
import torchvision

# An instance of your model.
model = torchvision.models.resnet18()

# An example input you would normally provide to your model's forward() method.
example = torch.rand(1, 3, 224, 224)

# Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
traced_script_module = torch.jit.trace(model, example)

trace된 ScriptModule은 이제 일반적인 파이토치 모듈로 사용가능

In[1]: output = traced_script_module(torch.ones(1, 3, 224, 224))
In[2]: output[0, :5]
Out[2]: tensor([-0.2698, -0.0381,  0.4023, -0.3010, -0.0448], grad_fn=<SliceBackward>)

Converting to Torch Script via Annotation

만약 모델이 특별한 형태나 흐름을 갖고있다면, Torch Script에 직접 기록하고 주석 달기를 원할 수 있다. 예를들어 모델이 아래와 같은 구조라면

import torch

class MyModule(torch.nn.Module):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

이 모듈의 forward 함수는 인풋에 따라 다른 흐름을 보이기 때문에 tracing에 알맞지 않음. 따라서 torch.jit.ScriptModule 에 부분집합을 만들고 @torch.jit.script_method 주석을 forward 메소드에 추가하는 방법을 ScriptModule로 컨버팅함

import torch

class MyModule(torch.jit.ScriptModule):
    def __init__(self, N, M):
        super(MyModule, self).__init__()
        self.weight = torch.nn.Parameter(torch.rand(N, M))

    @torch.jit.script_method
    def forward(self, input):
        if input.sum() > 0:
          output = self.weight.mv(input)
        else:
          output = self.weight + input
        return output

my_script_module = MyModule()

MyModule 오브젝트를 만들면 바로 ScriptModule 인스턴스를 만들어줌 

STEP 2: Script Module을 파일로 쓰기

아래의 명령어를 이용해 파일로 저장하면, 파이썬에 의존하지안고 C++에서 모델을 사용 가능함

traced_script_module.save("model.pt")

파이썬파트 여기까지가 끝

STEP 3: C++에서 Script Module 불러오기

LibTorch라는 파이토치 C++ API를 이용해서 불러와야 함. CMake와 LibTorch를 이용해서 모델을 불러와 이용하는 간단한 예제를 만들것

모델 불러오는 코드

#include <torch/script.h> // One-stop header.

#include <iostream>
#include <memory>

int main(int argc, const char* argv[]) {
  if (argc != 2) {
    std::cerr << "usage: example-app <path-to-exported-script-module>\n";
    return -1;
  }

  // Deserialize the ScriptModule from a file using torch::jit::load().
  std::shared_ptr<torch::jit::script::Module> module = torch::jit::load(argv[1]);

  assert(module != nullptr);
  std::cout << "ok\n";
}

파일 불러와 sharedpoint로 받고(C++에서의 torch.jit.ScriptModule) null 포인터인지 확인하는 코드임

Depending on LibTorch and Building the Application

위 코드가 example-app.cpp라는 스크립트에 적혀있다고 할때 CMakeLists.txt는 다음과 같음

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(Torch REQUIRED)

add_executable(example-app example-app.cpp)
target_link_libraries(example-app "${TORCH_LIBRARIES}")
set_property(TARGET example-app PROPERTY CXX_STANDARD 11)

마지막으로 필요한것은 LibTorch distribution이다.  download page 에서 최신버전 다운 가능함

받아서 압축 풀면 아래와 같음

libtorch/
  bin/
  include/
  lib/
  share/

lib/ 폴더는 공유 라이브러리 포함

include/ 폴더는 헤더파일 포함

share/ 폴더는 위에 find_package 명령어를 실행하기위한 필요한 CMake configuration을 포함

예제 경로가 다음과 같을때:

example-app/
  CMakeLists.txt
  example-app.cpp

다음의 명령어들을 실행하여 애플리케이션을 빌드함

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

/path/to/libtorch 는 LibTorch distribution을 압축 푼 풀 경로여야함. 만약 잘 됐다면 아래와 같아야함

root@4b5a67132e81:/example-app# mkdir build
root@4b5a67132e81:/example-app# cd build
root@4b5a67132e81:/example-app/build# cmake -DCMAKE_PREFIX_PATH=/path/to/libtorch ..
-- The C compiler identification is GNU 5.4.0
-- The CXX compiler identification is GNU 5.4.0
-- Check for working C compiler: /usr/bin/cc
-- Check for working C compiler: /usr/bin/cc -- works
-- Detecting C compiler ABI info
-- Detecting C compiler ABI info - done
-- Detecting C compile features
-- Detecting C compile features - done
-- Check for working CXX compiler: /usr/bin/c++
-- Check for working CXX compiler: /usr/bin/c++ -- works
-- Detecting CXX compiler ABI info
-- Detecting CXX compiler ABI info - done
-- Detecting CXX compile features
-- Detecting CXX compile features - done
-- Looking for pthread.h
-- Looking for pthread.h - found
-- Looking for pthread_create
-- Looking for pthread_create - not found
-- Looking for pthread_create in pthreads
-- Looking for pthread_create in pthreads - not found
-- Looking for pthread_create in pthread
-- Looking for pthread_create in pthread - found
-- Found Threads: TRUE
-- Configuring done
-- Generating done
-- Build files have been written to: /example-app/build
root@4b5a67132e81:/example-app/build# make
Scanning dependencies of target example-app
[ 50%] Building CXX object CMakeFiles/example-app.dir/example-app.cpp.o
[100%] Linking CXX executable example-app
[100%] Built target example-app

ResNet18모델을 성공적으로 만들었다면 아래와 같은 메세지가 뜸

root@4b5a67132e81:/example-app/build# ./example-app model.pt
ok

STEP 4: Executing the script module in C++

명령어 조금만 넣으면 바로 실행 가능 아래의 코드를 main() 함수에 추가

// Create a vector of inputs.
std::vector<torch::jit::IValue> inputs;
inputs.push_back(torch::ones({1, 3, 224, 224}));

// Execute the model and turn its output into a tensor.
at::Tensor output = module->forward(inputs).toTensor();

std::cout << output.slice(/*dim=*/1, /*start=*/0, /*end=*/5) << '\n';

모델의 인풋과 아웃풋은 IValue 타입이다 이후 forward 를 실행하고 텐서로 변환함

마지막줄은, 아웃풋의 처음 5개만 출력하는 것