도커를 활용한 TensorRT 환경 세팅하기

TensorRT 란?

TensorRT에 대해 간단하게 설명드리면 TensorRT는 NVIDIA platform에서 최적의 Inference 성능을 낼 수 있도록 Network compression, Network optimization 그리고 GPU 최적화 기술들을 대상 Deep Learning 모델에 자동으로 적용합니다.

딥러닝 모델을 GPU로 서비스를 할 때, TensorRT(https://developer.nvidia.com/tensorrt)로 모델을 변환하면 속도가 훨씬 빨라집니다.

그러면 물리 서버를 줄일 수 있어서, 운영 비용을 훨씬 절감할 수 있습니다.

 

TensorRT (https://developer.nvidia.com/tensorrt)

TensorRT가 모델을 최적화 할 때 사용하는 방법은 위 그림에 언급된 Layer & Tensor Fusion, Precision Calibration, Kernel Auto-Tuning, Dynamic Tensor Memory, Multi-Stream Execution과 같은 방법들 입니다.

 

TensorRT 환경 세팅

TensorRT 환경을 가장 간단하게 세팅할 수 있는 방법으로 소개 드리려고 합니다.

가장 간단하게 세팅할 수 있는 방법은 Docker를 활용하는 방법입니다.

Docker version은 19.03 이후 버전을 권장하고, nvidia-docker2도 설치되어 있어야 합니다.

TensorRT가 설치되어 있는 Docker image 리스트는 아래의 링크에서 확인할 수 있습니다.

https://ngc.nvidia.com/catalog/containers/nvidia:tensorrt

NVIDIA NGC

위의 링크에서 Docker image를 다운받아서 활용하시면 되는데, 여기서 주의하실 점이 하나 있습니다.

Host의 Nvidia driver 버전에 따라서 활용 할 수 없는 Docker image가 있습니다.

예를 들면, nvcr.io/nvidia/tensorrt:20.07.1-py3 이미지는 nvidia driver의 version이 450 이상이여야 하는 조건이 있습니다.

그 부분에 대해서는 다음의 첨부된 링크에서 확인하실 수 있습니다. https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html

 

제가 테스트한 버전은 Nvidia driver 440.100 버전이고, 활용한 Docker image는 nvcr.io/nvidia/tensorrt:20.03-py3 입니다.

먼저 아래의 명령어를 활용하여, Docker image를 받아줍니다. 

docker pull nvcr.io/nvidia/tensorrt:20.03-py3

그 다음, 아래의 명령어를 통해 컨테이너를 생성합니다.

docker run --gpus all -it --name trt_test nvcr.io/nvidia/tensorrt:20.03-py3

그러면 이제 간단하게 TensorRT 세팅이 모두 끝났습니다. 

이제 Pytorch나 Tensorflow와 같은 라이브러리를 설치하여 사용하시면 됩니다.

 

TensorRT 적용 테스트 (torch2trt)

이제 테스트도 해봐야겠죠?

TensorRT로 모델을 변환하는 것도 쉽지 않은 과정이지만, 간단한 예시로 자동으로 해주는 라이브러리가 있습니다.

TensorRT가 과연 제대로 세팅 되었는지, 한번 테스트를 해보겠습니다.

필요한 라이브러리는 torch, torchvision, torch2trt 입니다.

torch2trt에 대한 자세한 내용은 첨부된 링크에서 확인하실 수 있습니다. https://github.com/NVIDIA-AI-IOT/torch2trt

먼저, 아래의 명령어로 torch, torchvision을 설치합니다.

pip install torch torchvision

그 다음, torch2trt는 다음과 같은 과정을 통해 설치를 합니다.

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

그러면 이제 torch2trt 안에 있는 test.sh를 실행하면, 벤치 마크 테스트를 할 수 있습니다.

 

활용 방법에 대한 간단한 예제는 아래와 같습니다.

import torch
from torch2trt import torch2trt
from torchvision.models.alexnet import alexnet

# create some regular pytorch model...
model = alexnet(pretrained=True).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])

# execute
y = model(x) # torch model
y_trt = model_trt(x) #tensorrt model

 

딥러닝 모델을 GPU로 서비스하게 되면, 운영 비용이 크게 증가하게 됩니다.

서비스를 최적화하기 위하여 Distillation, Pruning과 같은 모델 네트워크를 수정하는 방법도 있지만, TensorRT와 같은 라이브러리를 추가로 활용만 해도 추론 속도를 높일 수 있습니다.

그래서, TensorRT는 설치가 꽤 까다로운데 이미 설정을 해둔 Docker image를 Nvidia에서 제공하고 있어서 오늘은 Docker를 활용하여 간단하게 TensorRT를 활용하는 방법을 알아보았으며, torch2trt를 통해 간단하게 테스트도 해보았습니다.

TensorRT를 더 잘 활용하기 위해서는 torch 모델을 바로 TensorRT로 변환하는 것 보다는 onnx를 거쳐서 변환하는 방법도 있습니다.