Pytorch를 TensorRT로 변환해서 사용하기

TensorRT 설치과정

TensoRT(trt)는 GPU inference 상황에서 최적의 optimization을 제공한다. 경우에 따라서는 거의 10배 이상 inference속도가 향상된다. 이 TensorRT를 사용하려면 조금 번거로운 설치과정을 해야한다. 그러나 설치만 잘 하면 변환 과정자체는 매우 쉽다.

1. CUDA 설치하기

cuda 버전 10.2 ~ 11.1 까지 지원된다. 여기서는 CUDA 11.0 버전을 사용한다.

2. PyCUDA 설치하기

pip install pycuda==2020.1

만약, CUDA를 새로 설치하거나 업데이트하면 pycuda역시 다시 재설치 or 업데이트가 필요하고 한다.

3. Pytorch 설치하기

pytorch는 1.5.0부터 테스트되었으나, 그전 버전에서도 동작할 수 있다고한다. 여기서는 pytorch 1.6버전을 사용한다.

주의할점

training 시스템과 inference 시스템이 같을 경우, 가능하면 동일한 cuDNN 버전을 사용하는 것이 좋다.

3. cuDNN 설치하기

8.0.4 version의 dev/runtime/example 의 세가지 deb 파일을 여기서 다운로드한다.

sudo dpkg -i libcudnn8_8.0.4.30-1+cuda11.0_amd64.deb

sudo dpkg -i libcudnn8-dev_8.0.4.30-1+cuda11.0_amd64.deb

sudo dpkg -i libcudnn8-samples_8.0.4.30-1+cuda11.0_amd64.deb

정상적으로 설치가 되었다면, 이 방법을 따라 cudnn_sample을 컴파일해서 실행해보자.

cuDNN version check 하는 방법

dpkg -l | grep cudnn

4. TensorRT 설치하기

여기서 내 Ubuntu 버전 및 CUDA 버전에 알맞는 TensorRT를 다운로드한다. 이 때 결국 제대로 tensorrt를 사용하려면 TAR과 DEB를 모두 설치해야한다. DEB 파일의 경우 ubuntu tensorrt 프로그램을 설치해주는 것이고, TAR file은 python tensorrt 모듈을 설치해준다.

TAR 설치하기

tar xvzf TensorRT-7.2.1.6.Ubuntu-18.04.x86_64-gnu.cuda-11.0.cudnn8.0.tar.gz

(이 방식의 경우, TensorRT 폴더가 위치한 곳이 library가 실제로 존재하는 곳이므로 적절한 경로를 잘 설정해주어야한다.)

mv TensorRT....  /usr/local/lib/

export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/lib/TensorRT-7.2.1.6/lib

pip install python/tensorrt-7.2.1.6-cp37-none-linux_x86_64.whl

pip install uff/uff-0.6.9-py2.py3-none-any.whl

pip install graphsurgeon/graphsurgeon-0.4.5-py2.py3-none-any.whl

pip install onnx_graphsurgeon/onnx_graphsurgeon-0.2.6-py2.py3-none-any.whl

설치확인

cd samples/samplesMNIST

make

cd ../../bin/

./sample_mnist

삭제하기

간단히 LD_LIBRARY_PATH 에 등록된 TensorRT 디렉토리를 그냥 지워버리면 된다.

DEB 설치하기

dpkg -i nv-tensorrt-repo-ubuntu1804-cuda11.0-trt7.2.1.6-ga-20201007_1-1_amd64.deb

apt-key add /var/nv-tensorrt-repo-cuda11.0-trt7.2.1.6-ga-20201007/7fa2af80.pub

apt-get update

apt-get install -y tensorrt

apt-get install -y python3-libnvinfer-dev

apt-get install -y onnx-graphsurgeon

설치확인

dpkg -l | grep TensorRT

삭제하기

sudo apt-get purge "libnvinfer*"

sudo apt-get purge graphsurgeon-tf onnx-graphsurgeon

5. Torch2trt 설치하기

pytorch 모델을 바로 trt로 변환해주는 python 프로젝트이다. 설치가 매우 간편하며, 사용성 또한 편리하다.

게다가 onnx를 거쳐서 porting하는 옵션도 있어서 매우 flexible하다.

설치

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

cd torch2trt

sudo python setup.py install --plugins

Python에서 trt호출하기

여기에 안내된 것과 같이 변환하면, 쉽게 inference할 수 있다.

onnx 거쳐서 porting하기

torch2trt(self.conf.model, [x], use_onnx=True)