torch-TensorRT 를 사용하며. torchscript 란?

회사에서 사용하는 로직을 수정할 일이 생겨서

nvidia의 https://developer.nvidia.com/ko-kr/blog/torch-tensorrt%EB%A5%BC-%ED%86%B5%ED%95%B4-pytorch%EC%97%90%EC%84%9C-%EC%B6%94%EB%A1%A0-%EC%86%8D%EB%8F%84-%EC%B5%9C%EB%8C%80-6%EB%B0%B0-%ED%96%A5%EC%83%81%ED%95%98%EA%B8%B0/

Torch-TensorRT를 통해 PyTorch에서 추론 속도 최대 6배 향상하기

를 적용하기 위해 torch-tensorRT를 사용하고 있다.

C++로 사용하려고 설치단에서도 고생을 좀 했다. 

Jetpack 기반에 설치한 거라 바이너리로 받아서 하나하나 설치..(왈칵)

 

어쨌든 돌아가는 것도 확인하고 이제 모델 추로을 하고 있는데 모델들을 torchscript로 변환하는 것이 먼저라고 해서

찾아보았다.

 

Pytorch는 공식 한국어 사이트(?) 가 있어서 그나마 낫다.

ndivia의 포럼은..^^..영어 그잡채

 

https://tutorials.pytorch.kr/recipes/torchscript_inference.html

TorchScript로 배포하기

 

pytorch도 결국 까보면 *.so 파일을 사용하는데, 어찌됐든 c++에서 사용하기 위해서는 모델을 토치 스크립트로 변환을 해야한다.

 

** 위의 사이트에 나온 예시

 

모델을 python에서 저장을 해서

import torch
import torch.nn.functional as F
import torchvision.models as models

r18 = models.resnet18(pretrained=True)       # 이제 사전 학습된 모델의 인스턴스가 있습니다.
r18_scripted = torch.jit.script(r18)         # *** 여기가 TorchScript로 내보내는 부분입니다.
dummy_input = torch.rand(1, 3, 224, 224)     # 빠르게 테스트 해봅니다.

r18_scripted.save('r18_scripted.pt') # 예시

c++파일을 만든다. 사이트 참고.

#include <torch/script.h>
#include <torch/nn/functional/activation.h>


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

    std::cout << "Loading model...\n";

    // ScriptModule을 역직렬화(deserialize) 합니다.
    torch::jit::script::Module module;
    try {
        module = torch::jit::load(argv[1]);
    } catch (const c10::Error& e) {
        std::cerr << "Error loading model\n";
        std::cerr << e.msg_without_backtrace();
        return -1;
    }

    std::cout << "Model loaded successfully\n";

    torch::NoGradGuard no_grad; // autograd가 꺼져있는지 확인합니다.
    module.eval(); // dropout과 학습 단의 레이어 및 함수들을 끕니다.

    // 입력 "이미지"를 생성합니다.
    std::vector<torch::jit::IValue> inputs;
    inputs.push_back(torch::rand({1, 3, 224, 224}));

    // 모델을 실행하고 출력 값을 tensor로 뽑아냅니다.
    at::Tensor output = module.forward(inputs).toTensor();

    namespace F = torch::nn::functional;
    at::Tensor output_sm = F::softmax(output, F::SoftmaxFuncOptions(1));
    std::tuple<at::Tensor, at::Tensor> top5_tensor = output_sm.topk(5);
    at::Tensor top5 = std::get<1>(top5_tensor);

    std::cout << top5[0] << "\n";

    std::cout << "\nDONE\n";
    return 0;
}

이 파일이 make하고 나면 우리가 저장한 ts파일을 불러오고, 이미지 tensor를 만들고 입력에 대한

추론을 실행한다.

 

CMakeLists.txt 파일 생성.

cmake_minimum_required(VERSION 3.0 FATAL_ERROR)
project(custom_ops)

find_package(Torch REQUIRED)

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

그리고 cmake

cmake -DCMAKE_PREFIX_PATH=<path to your libtorch installation>
make

 

위에서 만든 c++ 파일, pt파일 경로 넣어주기

$ ./ts-infer r18_scripted.pt

그럼 나온다!

 

 

토치스크립트로 변환된 파일이 있으면 torchvision 종속성이 없다고한다.

학습 가중치와 연산 그래프가 같이 저장된다.