Docker를 이용한 torch-tensorrt 설치하기(with. ubuntu)

까람·2023년 1월 3일

AWS EC2 인스턴스를 대여해서 진행되었으며, Ubuntu 20.04 LTS, V100 GPU를 기준으로 작성되었습니다.

TensorRT를 사용해서 추론을 가속화하기 위해서는 여러 방법이 있는 것 같지만, 많은 분들이 도커를 활용하는 것이 가장 편했다고들 하시니, 저도 도커를 활용해서 설치해보겠습니다.

우선 NIDIA GPU 드라이버를 설치해보겠습니다.

NVIDIA GPU Driver 설치

본인의 GPU와 OS 환경에 적합한 드라이버를 찾아 설치해야합니다.

NVIDIA Driver 다운로드 페이지

pytorch의 stable 최신 버전이 cuda 11.7까지 지원하는 하는 것으로 알고 있어서 CUDA Toolkit은 11.7 버전으로 했습니다.

아래 명령어로 다운로드 받은 설치 파일을 실행해주었습니다. sudo 명령어로 하지 않으면 설치 안되더라구요.

$ sudo apt update
$ sudo apt upgrade
$ sudo sh NVIDIA-Linux-x86_64-515.86.01.run

도중에 아래와 같은 에러가 발생해서 gcc를 다운로드 받는 것으로 해결했습니다.

 ERROR: Unable to find the development tool `cc` in your path; please make sure that you have the package 'gcc' installed.  If gcc is installed on your system, then
         please check that `cc` is in your PATH.
$ sudo apt install gcc

이번엔 make 어쩌고 에러가 떠서 이또한 다운로드 해주었습니다.

  ERROR: Unable to find the development tool `make` in your path; please make sure that you have the package 'make' installed.  If make is installed on your system, then
         please check that `make` is in your PATH.
$ sudo apt install make

이후 아래와 같은게 떳는데 무시했고, 32 bit 라이브러리 다운받겠냐는 것도 안받았습니다.

  WARNING: nvidia-installer was forced to guess the X library path '/usr/lib64' and X module path '/usr/lib64/xorg/modules'; these paths were not queryable from the
           system.  If X fails to find the NVIDIA X driver module, please install the `pkg-config` utility and the X.Org SDK/development package for your distribution
           and reinstall the driver.

gcc, make 패키지 설치 후 무사히 GPU 드라이버 설치가 되었습니다.

드라이버가 잘 설치됬는지 확인해보았고, 515 버전으로 잘 설치가 된 것 같습니다.

$ nvidia-smi

Docker 설치

공식 도커 인스톨 가이드 페이지를 참고해서 진행했습니다.

우선 도커 엔진을 다운받기 전에 레포지토리 셋업을 해야한다고 하네요.

$ sudo apt-get install \
    ca-certificates \
    curl \
    gnupg \
    lsb-release

도커 공식 GPG key를 추가해주겠습니다.(GPG key가 뭔지 모름 그냥 무지성 복붙)

sudo mkdir -p /etc/apt/keyrings
curl -fsSL https://download.docker.com/linux/ubuntu/gpg | sudo gpg --dearmor -o /etc/apt/keyrings/docker.gpg

아래 명령어로 레포지토리를 셋업 해주라고 하네요.

$ echo \
  "deb [arch=$(dpkg --print-architecture) signed-by=/etc/apt/keyrings/docker.gpg] https://download.docker.com/linux/ubuntu \
  $(lsb_release -cs) stable" | sudo tee /etc/apt/sources.list.d/docker.list > /dev/null

이제 도커 엔진을 다운로드 받습니다.

$ sudo apt-get update
$ sudo apt-get install docker-ce docker-ce-cli containerd.io docker-compose-plugin

잘 설치가 되었는지 샘플 이미지로 확인해봅니다.

$ sudo docker run hello-world
Unable to find image 'hello-world:latest' locally
latest: Pulling from library/hello-world
2db29710123e: Pull complete
Digest: sha256:c77be1d3a47d0caf71a82dd893ee61ce01f32fc758031a6ec4cf1389248bb833
Status: Downloaded newer image for hello-world:latest

Hello from Docker!
This message shows that your installation appears to be working correctly.

To generate this message, Docker took the following steps:
 1. The Docker client contacted the Docker daemon.
 2. The Docker daemon pulled the "hello-world" image from the Docker Hub.
    (amd64)
 3. The Docker daemon created a new container from that image which runs the
    executable that produces the output you are currently reading.
 4. The Docker daemon streamed that output to the Docker client, which sent it
    to your terminal.

To try something more ambitious, you can run an Ubuntu container with:
 $ docker run -it ubuntu bash

Share images, automate workflows, and more with a free Docker ID:
 https://hub.docker.com/

For more examples and ideas, visit:
 https://docs.docker.com/get-started/

설치된 버전도 확인해봅니다.

$ docker --version
Docker version 20.10.22, build 3a2c30b

설치가 잘된 것 같습니다.

torch-tensorrt 세팅

포스팅을 하고 있는 지금은 23년 1월 3일인데요. 22년 6월 16일경 NVIDIA가 Torch와 함께 공식적으로 손을 잡고, TensorRT가 공식적으로 Pytorch에 도입되었나 봅니다. Torch-TensorRT의 깃허브가 공식 레포가 된 것 같습니다. 설치, 활용 가이드 등 문서도 공식 Pytorch 사이트에서 업데이트가 된다고 하네요.

저의 입장에서 좋은건 기존에는 TensorRT 환경 도커 이미지까지는 NVIDA에서 제공해주고 있는 것이 있었는데, 여기서 추가적으로 사용하는 딥러닝 프레임워크에 따라서 의존성을 따로 갖추는 작업이 필요했습니다. 근데 이제 Torch + CUDA Toolkit + cuDNN + TensorRT + Torch-Tensorrt까지 모든 의존성과 호환성이 다 갖추어져 있는 이미지를 제공해주는 듯 합니다...(드라이버는 호스트꺼를 빌려쓰는 형태인 것 같아 따로 설치해주어야 할듯요)

호환성을 고려한 환경을 세팅하는 부분이 가장 어려웠는데, 그 어려운걸 쉽게 쉽게 갈 수 있게 됩니다.
환경은 다 갖추어져 있으니 컨테이너 띄워서 진짜 코드만 넣고, 돌려보면 되는 수준...

세팅 과정은 PyTorch x NVIDIA 설치 가이드 문서를 따르겠습니다.

22.12-py3 버전이 최신버전의 이미지 진행했습니다.

-v 옵션을 통해 호스트 프로젝트 디렉토리 하나를 마운트 해주었습니다.

$ docker run --gpus all -it --ulimit memlock=-1 --ulimit stack=-1 --ulimit core=-1 --ipc=host --shm-size=32gb --name tensorrt --rm -v /test/test-project/:/workspace/test-project nvcr.io/nvidia/pytorch:22.12-py3

이미지가 다 받아지면, 컨테이너에 접속이 됩니다. 실제로 GPU 연동이 되었는지 확인해보겠습니다.

$ nvidia-smi
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.86.01    Driver Version: 515.86.01    CUDA Version: 11.8     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla V100-SXM2...  Off  | 00000000:00:1E.0 Off |                    0 |
| N/A   38C    P0    40W / 300W |      0MiB / 16384MiB |      1%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|  No running processes found                                                 |
+-----------------------------------------------------------------------------+

드라이버는 잘 인식하고 있는 것 같습니다. 그리고 CUDA Toolkit도 잘 설치되어있는지 확인해보겠습니다.

$ nvcc -V
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2022 NVIDIA Corporation
Built on Wed_Sep_21_10:33:58_PDT_2022
Cuda compilation tools, release 11.8, V11.8.89
Build cuda_11.8.r11.8/compiler.31833905_0

CUDA Toolkit도 11.8로 잘 설치가 되어있네요. cuDNN도 확인해보겠습니다.

$ cat /usr/include/x86_64-linux-gnu/cudnn_v*.h | grep CUDNN_MAJOR -A 2
#define CUDNN_MAJOR 8
#define CUDNN_MINOR 7
#define CUDNN_PATCHLEVEL 0
--
#define CUDNN_VERSION (CUDNN_MAJOR * 1000 + CUDNN_MINOR * 100 + CUDNN_PATCHLEVEL)

/* cannot use constexpr here since this is a C-only file */

cuDNN도 8.7.0 버전으로 잘 설치가 되어있네요. 이제 마지막으로 torch 라이브러리단에서 GPU를 사용할 수 있는지 확인해보겠습니다.

$ python
Python 3.8.10 (default, Nov 14 2022, 12:59:47)
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import torch
>>> print(torch.cuda.is_available())
True

torch도 잘 동작하는걸 알 수 있습니다. 이제 tensorrt를 적용해서 확인해보겠습니다. 아래 테스트 코드는 다른 분의 코드를 참고했습니다.

import torch
from torchvision.models.mobilenet import mobilenet_v2
import torch_tensorrt


net = mobilenet_v2(pretrained=True).cuda().eval()

trt_ts_module = torch_tensorrt.compile(net,
    inputs = [torch_tensorrt.Input( # Specify input object with shape and dtype
            min_shape=[1, 3, 112, 112],
            opt_shape=[1, 3, 224, 224],
            max_shape=[1, 3, 448, 448],
            # For static size shape=[1, 3, 224, 224]
            dtype=torch.half) # Datatype of input tensor. Allowed options torch.(float|half|int32|bool)
    ],
    enabled_precisions = {torch.half}, # Run with FP16
)
torch.jit.save(trt_ts_module, "trt_torchscript_module.ts") # save the TRT embedded Torchscript
import time
import os


def get_model_size(net):
    torch.save(net.state_dict(), "tmp.pth")
    model_size = os.path.getsize("tmp.pth") / 1e6
    os.remove("tmp.pth")
    return model_size


float32_data = torch.ones([1, 3, 224, 224]).cuda()
float16_data = torch.ones([1, 3, 224, 224], dtype=torch.half).cuda()
start = time.time()
for i in range(1000):
    result = net(float32_data) # run inference
end = time.time()
print(f"Pytorch time cost: {end - start:.4f}ms, Mode Size: {get_model_size(net):.4f}MB")


start = time.time()
for i in range(1000):
    result = trt_ts_module(float16_data) # run inference
end = time.time()
print(f"TensorRT time cost: {end - start:.4f}ms, Model Size: {get_model_size(trt_ts_module):.4f}MB")

아래는 테스트 결과입니다. 추론 성능과 모델 사이즈에서 엄청나게 큰 차이를 보이는 것으로 보아 TensorRT가 잘 동작한다는 것을 알수 있습니다.

Pytorch time cost: 8.8832ms, Mode Size: 14.2512MB
TensorRT time cost: 0.4586ms, Model Size: 0.0005MB

끝.

profile
한 줄 소개 쓸 자랑거리가 없다..

0개의 댓글