[TensorRT 튜토리얼] 5. torch2trt

TensorRT를 사용하는 것은 어렵다. 아무리 ONNX, python API의 도움을 받아도 직접 conversion, deployment를 깔끔하게 해내긴 어렵다. 

 

https://github.com/NVIDIA-AI-IOT/torch2trt

GitHub - NVIDIA-AI-IOT/torch2trt: An easy to use PyTorch to TensorRT converter

torch2trt는 TensorRT 사용을 고민하는 사람들이 가장 흔하게 사용할 조합인 pytorch model 변환 + python inference를 손쉽게 하기 위해 제작된 일종의 toolbox다. 직접 conversion, deployment 코드를 API를 사용해서 구현할 필요없이 torch2trt 호출만으로 모델 변환을 끝내고 동작 가능한 형태로 바꿔준다. 

 

@tensorrt_converter 데코레이터를 통해 custom layer를 변환하는 기능을 지원하기 때문에 사용성이 꽤 좋은 편이다.

 

Prerequiste

당연하게도 tensorrt가 설치되어 있어야 한다. python API만 사용하는 toolbox기 때문에 python wheel로 설치해도 무방하다. 

 

설치는 이전 글 [TensorRT 튜토리얼] 1. Installing TensorRT 참고 

 

Install

git clone https://github.com/NVIDIA-AI-IOT/torch2trt
cd torch2trt
python setup.py install

 

마지막 install 명령에서 다음과 같은 오류가 뜰 수도 있다.

running install
error: can't create or remove files in install directory

The following error occurred while trying to add or remove files in the
installation directory:

    [Errno 13] Permission denied: '/usr/local/lib/python3.8/dist-packages/test-easy-install-130214.write-test'

 

permission 문제라고 해서 sudo 붙이면 ModuleNotFoundError : No module named 'tensorrt'  추가 오류를 만날 수 있다.

python3 setup.py install --user

 

마지막에 --user 를 추가해주면 설치 완료된다.

 

cmake -B build . 
cmake --build build --target install 
ldconfig

 

추가로 torch2trt plugins library를 선택적으로 설치할 수 있는데, 쓸지 안 쓸지 모르지만 설치는 무조건 해두자.

CMake Error at cmake_install.cmake:52 (file):
  file INSTALL cannot copy file
  "/home/jseob/Desktop/yjs/codes/torch2trt/build/libtorch2trt_plugins.so" to
  "/usr/local/lib/libtorch2trt_plugins.so": Permission denied.


make: *** [Makefile:100: install] Error 1

 

또 permision denied 오류를 만날 수 있는데, 여기선 sudo 쓰면 된다. 

cmake -B build .
sudo cmake --build build --target install
sudo ldconfig

 

import tensorrt as trt
import torch2trt

import가 제대로 되는 것만 확인하면 설치 완료!

 

Usage

import torch
import torchvision.models as models
from torch2trt import torch2trt, TRTModule

### create some regular pytorch model...
model = models.resnet50(weights=True, progress=False).eval().cuda()
### create example data
x = torch.ones([1, 3, 224, 224]).cuda()

### convert to TensorRT feeding sample data as input
model_trt = torch2trt(model, [x])

### save and load
# torch.save(model_trt.state_dict(), SAVE_PATH)
# model_trt = TRTModule
# model_trt.load_state_dict(torch.load(SAVE_PATH))

y = model(x)
y_trt = model_trt(x)

### check the output against PyTorch
print(torch.max(torch.abs(y - y_trt)))

# tensor(0.0081, device='cuda:0', grad_fn=<MaxBackward1>)

 

사용도 pytorch 문법과 매우 유사하고 pytorch tensor를 지원하기 때문에 아주 편하다. 

 

Process finished with exit code 139 (interrupted by signal 11: SIGSEGV)

 

위와 같은 오류가 뜰 때가 있었는데 torch2trt.py line 694, outputs=module(*inputs) 부분을 못 넘어가고 segmentation fault가 떴다. 

 

원인은 TensorRT를 설치할 때 사용했던 CUDA 버전과 맞지 않는 pytorch를 사용했기 때문이다. 예를 들어, torch2.0.1+cu118을 CUDA 12.1 사용해는 TensorRT와 사용할 경우 오류가 난다. 이 경우 간단하게 pytorch 버전을 CUDA에 맞게 업데이트 해주면 된다.