Pytorch를 TensorRT로 변환하기 위한 환경 설정

1. Introduction

TensorRT는 학습된 딥러닝 모델을 최적화하여 NVIDIA GPU 상에서 추론하는 속도를 수십 ~ 수십배까지 향상시켜주는 모델 최적화 엔진입니다. 지원하는 프레임워크는 pytorch, tensorflow, keras, onnx 등이 있습니다. 여러 프레임워크에서 직접적인 TensorRT 변환을 지원하고 있지만 다소 제약이 있으며, 대신 onnx로의 변환은 대부분의 프레임워크에서 더 나은 지원을 해주고 있으며 onnx로 변환한 모델은 다시 TensorRT로 변환이 가능합니다. 저희는 대부분의 연구/개발을 pytorch로 하기 때문에 pytorch을 onnx로 변환하고 이를 TensorRT로 추론하는 방법을 선택하였습니다.

이번 글에서는 TensorRT로 변환했을 때 속도가 얼마나 빨라지는지를 비교하는 모델로 Google에서 발표한 Bert모델을 사용합니다. 이를 이용해 pytorch모델을 onnx로 변환하고 onnx모델을 다시 TensorRT로 변환하기 위한 환경설정을 Docker로 하는 방법을 소개합니다.

2. Version Matching

TensorRT를 설치하는데 있어 주의해야 할 점 중 하나는 Ubuntu, cuda, pytorch, python 등등 모듈들의 버전을 잘 맞춰줘야 한다는 점입니다. 이번 글에서는 아래와 같은 버전을 사용했습니다.

ubuntu==18.04
cuda==10.0
cudnn==7
python==3.7
pytorch==1.4.0
torchvision=0.4.0
onnx==1.7
transformers==3.0.2

3. Download TensorRT

이들을 설치하기 이전에 가장 중요한 TensorRT의 압축파일을 NVIDIA 홈페이지에서 다운로드합니다. Version Matching에서 설명했듯이 Ubuntu 18.04, CUDA 10.0을 사용할 것이기 때문에 아래 그림과 같은 tar package을 다운로드 합니다.

다운로드 받은 압축파일을 TensorRT-7.0.0.11 디렉토리에 압축풀기를 합니다.

4. Dockerfile

Dockerfile을 구성합니다.
가장 먼저 OS에 관한 명령어를 넣습니다. Version Matching에서 설명하였듯이 Ubuntu 18.04, cuda 10.0, cudnn 7에 맞는 환경을 설치합니다. 물론 직접 설치하는 방법을 사용할 수 있지만 nvidia에서 docker image로 제공하고 있기 때문에 이를 사용하도록 하겠습니다.

FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04

Python3.7을 설치합니다.

RUN apt-get -y update
RUN apt-get -y upgrade
RUN apt-get -y dist-upgrade
RUN apt-get -y install build-essential python-dev python-setuptools python-pip python-smbus
RUN apt-get -y install libncursesw5-dev libgdbm-dev libc6-dev
RUN apt-get -y install zlib1g-dev libsqlite3-dev
RUN apt-get -y install libssl-dev openssl
RUN apt-get -y install libffi-dev wget
RUN wget https://www.python.org/ftp/python/3.7.4/Python-3.7.4.tgz
RUN tar xvfz Python-3.7.4.tgz
RUN cd Python-3.7.4 && ./configure && make && make install

다음은 딥러닝 프레임워크와 관련 패키지를 설치합니다.

RUN pip3 install tensorflow-gpu==1.14.0
RUN pip3 install torch==1.4.0
RUN pip3 install torchvision==0.4.0
RUN pip3 install transformers==3.0.2
RUN pip3 install onnx==1.7.0
RUN pip3 install onnx-tf
RUN pip3 install pycuda

마지막으로 압축 해제한 TensorRT 디렉토리 내부의 파일들을 이용하여 TensorRT를 설치합니다.

COPY TensorRT-7.0.0.11 /TensorRT-7.0.0.11
RUN cd TensorRT-7.0.0.11/python && pip3 install tensorrt-7.0.0.11-cp37-none-linux_x86_64.whl
RUN cd TensorRT-7.0.0.11/uff && pip3 install uff-0.6.5-py2.py3-none-any.whl
RUN cd TensorRT-7.0.0.11/graphsurgeon && pip3 install graphsurgeon-0.4.1-py2.py3-none-any.whl
ENV LD_LIBRARY_PATH "$LD_LIBRARY_PATH:/usr/local/cuda-10.0/lib64:/usr/local/cuda/extras/CUPTI/lib64:/TensorRT-7.0.0.11/lib"

모두 병합한 Dockerfile은 아래와 같이 구성됩니다.

FROM nvidia/cuda:10.0-cudnn7-devel-ubuntu18.04

RUN apt-get -y update
RUN apt-get -y upgrade
RUN apt-get -y dist-upgrade
RUN apt-get -y install build-essential python-dev python-setuptools python-pip python-smbus
RUN apt-get -y install libncursesw5-dev libgdbm-dev libc6-dev
RUN apt-get -y install zlib1g-dev libsqlite3-dev
RUN apt-get -y install libssl-dev openssl
RUN apt-get -y install libffi-dev wget
RUN wget https://www.python.org/ftp/python/3.7.4/Python-3.7.4.tgz
RUN tar xvfz Python-3.7.4.tgz
RUN cd Python-3.7.4 && ./configure && make && make install

RUN pip3 install tensorflow-gpu==1.14.0
RUN pip3 install torch==1.4.0
RUN pip3 install torchvision==0.4.0
RUN pip3 install transformers==3.0.2
RUN pip3 install onnx==1.7.0
RUN pip3 install onnx-tf
RUN pip3 install pycuda

COPY TensorRT-7.0.0.11 /TensorRT-7.0.0.11
RUN cd TensorRT-7.0.0.11/python && pip3 install tensorrt-7.0.0.11-cp37-none-linux_x86_64.whl
RUN cd TensorRT-7.0.0.11/uff && pip3 install uff-0.6.5-py2.py3-none-any.whl
RUN cd TensorRT-7.0.0.11/graphsurgeon && pip3 install graphsurgeon-0.4.1-py2.py3-none-any.whl
ENV LD_LIBRARY_PATH "$LD_LIBRARY_PATH:/usr/local/cuda-10.0/lib64:/usr/local/cuda/extras/CUPTI/lib64:/TensorRT-7.0.0.11/lib"

WORKDIR /app

5. Docker build

Dockerfile을 빌드하기 전에 주의해야 할 점이 있습니다. 이전에 압축 해제한 TensorRT 디렉토리를 Dockerfile에서 읽을 수 있는 위치로 옮겨야 합니다. 압축 해제한 TensorRT 디렉토리를 Dockerfile과 같은 위치에 옮기면 Dockerfile을 빌드하면서 읽을 수 있습니다.

└─xxx
   ├─Dockerfile
   └─TensorRT-7.0.0.11

위의 디렉토리를 확인한 후 Docker build를 진행합니다.

docker build -t tensorrt:latest .

6. Conclusion

몇 가지 명령어로 설치가 정확히 되었는지 확인할 수 있습니다.

$ NV_GPU=0 nvidia-docker run -it --rm --shm-size 16G tensorrt:ckg which convert-to-uff
/usr/local/bin/convert-to-uff

NV_GPU=0 nvidia-docker run -it --rm --shm-size 16G tensorrt:ckg python3 -c "import tensorrt"

환경 설정에 대한 내용은 여기까지이며 TensorRT를 이용하여 추론 최적화를 진행하는 내용은 다음 글을 참조해주시길 바랍니다.

[##차금강##]