update Nov.23.22: 오류수정

Prologue

지난주 즈음부터 나에게 주어진 업무.

"TensorRT가 뭔지, native model, cuda engine model을 각각 비교해보고, cuda engine으로 변환이 안 되는 레이어는 어떻게 하는지 알아보고 발표하세요"

Envrionment

일단 예시코드부터 돌려본다. 아니, 그 전에 세팅부터. 예전에는 파일을 받아와서 직접 빌드했던 것 같다. 요즘은 그런 거 안 해도 된다. 엔비디아에서 관리하는 컨테이너가 있고 거의 매달 업데이트한다.[1]

$ docker pull nvcr.io/nvidia/tensorrt:<YY.mm>-py<x>

프로젝트에 필요한 프레임웍, 라이브러리버전에 맞는 걸로 골라서 쓰면 된다.

이렇게 받은 이미지를 활용하면 TensorRT를 쓸 수 있다. 컨테이너에는 TensorRT말고는 아무것도 없어서 나머지 필요한 건 requirments.txt를 만들어뒀다가 한꺼번에 설치하면 된다.

그러면 오늘은 ResNet50을 cuda engine으로 만들어볼거다. 참고로 내가 쓰고 있는 TensorRT버전은 8.2.5인데 웹상에서 돌아다니는 예시코드들은 7.x 혹은 그 이전이라 쓸 수 가 없어서 docs를 보고 일일이 찾아야 했다.

native model

먼저 ResNet50부터 불러와보자.

import torch
from torchvision import models, transforms

resnet50 = models.resnet50(weights='DEFAULT', progress=False).eval()

이렇게 하면 모델개발 없이 imagenet을 학습한 모델을 불러올 수 있다. 그러면 아무 이미지 하나 골라와서 어떻게 예측하는지 보자.

import numpy as np
from skimage import io
from skimage.transform import resize

url = 'https://url/of/the/image'
img = io.imread(url)
img = resize(img, (224, 224))
img = np.asarray(img)

# Normalize
img -= np.array([123.68, 116.779, 103.939])
img /= np.array([58.393, 57.12, 57.375])

# [Batch, C, H, W]
img = torch.from_numpy(img).unsqueez(0).permute(0, 2, 3, 1).cuda()

resnet50_gpu = models.resnet50(weights='DEFAULT', progress=False).cuda().eval()

# Warmup
with torch.no_grad():
    resnet50_gpu(img)

with torch.no_grad():
    tic = time.time()
    out = resnet50_gpu(img)
print(time.time() - tic)
print(torch.argmax(torch.sigmoid(out)))

내가 가진 GPU에서는 이미지 한 장을 추론하는데 0.009초가 걸리고 207번 index를 가리킨다. 이제 이걸 cuda engine으로 바꿀 거다. torch2trt를 쓰면 바로 할 수 있는 것처럼 보인다. 그렇지만 좀더 범용으로 쓸 수 있는 ONNX를 쓸 거다. 그러니까 전체적인 과정은 native model $\to$ ONNX model $\to$ cuda engine 이렇게 3단계를 거친다.

torch2onnx

import onnx

# convert torch native model to ONNX model
batch_size = 1
dummy_input = torch.randn((batch_size, 3, 224, 224), dtype = torch.float32)
torch.onnx.export(resnet50, dummy_input, 'path/to/onnx/file', verbose=False)

여기에서 특이하다고 생각했던 건 inference 모드에 있는 모델과 함께 입력크기를 정해줘야 한다는 점이다. 이렇게 해야 각 레이어의 입, 출력크기를 알 수 있어서 그러는 것 같다. 이건 아니고 torch.onnx.export를 실행하면 내부적으로 torch.jit.trace가 실행된다. 이건 모델객체와 인풋더미를 요구하는데 모델 객체에다가 인풋을 넣으면서 각 레이어가 가진 계산 그래프를 추적한다. 커스텀 레이어의 경우 간혹 입력값에 의존해서 계산방법이 달라지는 레이어가 있는데 이 경우도 이야기 하려면 글이 너무 길어지므로 넘어가기로 한다.

onnx2trt

편하게는 trtexec 명령어를 이용하면 편하게 최적화할 수 있다. 그렇지만 좀더 빡빡하게 하려면 python이든, c++든 스크립트를 짜야 한다. 이건 시도해봤는데 중간에 뭐가 잘못됐는지 성능이 많이 떨어져서 이번에는 trtexec을 이용하기로 한다.

$ trtexec --onnx=path/to/onnx/model \
		--saveEngine=path/to/cuda/engine/file \
        --inputIOFormats=fp16:chw \
        --outputIOFormats=fp16:chw \
        --fp16

이 명령어가 지정한 경로에 cuda engine이라고 부르는 binary 파일을 만들어준다.

trt inference

앞에서 만들었던 cuda engine은 weight 파일처럼 바로 쓸 수 있는 게 아니라 읽어와서 변환을 한 번 거쳐야 한다. 이걸 deserialize라고 부른다. 반대로 앞에서 cuda engine을 만드는 과정을 serialize라고 한다.

import tensorrt as trt
import pycuda.drive as cuda
import pycuda.autoinit

runtime = trt.Runtime(logger)
f = open('path/to/cuda/engine/file' as 'rb')
engine = runtime.deserialize_cuda_engine(f.read())
context = engine.create_execution_context()
stream = cuda.Stream()

img = np.asarray(img, dtype = np.float16)
output = np.empty([batch_size, 1000], dtype = np.float16)

d_input = cuda.mem_alloc(1 * img.nbytes)
d_output = cuda.mem_alloc(1 * output.nbytes)
bindings = [int(d_input), int(d_output)]

def inference(inputs, outputs, bindings, stream, context):
    cuda.memcpy_htod(bindings[0], inputs, stream)
    context.execute_(bindings, stream.handle, None)
    cuda.memcpy_dtoh(outputs, bindings[1], stream)
    stream.synchronize()
    return outputs

# Warmup
inference(img, output, bindings, stream, context)

tic = time.time()
out = inference(img, output, bindings, stream, context)
print(time.time() - tic)
print(torch.argmax(torch.sigmoid(torch.from_numpy(out))))

이건 대략 0.003초가 걸리고 inference결과는 native model과 같다. cuda engine model이 3배정도 빠르다. 물론 정밀도를 낮춰서 그렇기도 하지만 그것 말고도 kernel fusion, model을 하드웨어에 최적화해줘서 그렇다. 진짜 기술은 명령어 한 줄을 치는 게 아니라 최적화 스크립트를 일일이 짜서 script.py를 실행하는 거지 않을까 싶다.

Epilogue

나는 분명 python 스크립트를 작성하고 있는데 중간에 memory allocation하는 코드때문에 그런가 묘하게 c프로그래밍하는 기분이 든다. 까다로운 점은 라이브러리 버전이 달라지면 이전버전에서 쓰던 class나 method가 없어지거나 아예 다른 class가 담당하기도 해서 온라인상에 돌아다니는 코드를 보는 것은 큰 도움이 되지는 않았다. 오히려 공식 repo나 docs, guide를 꾸역꾸역 읽는 게 더 나았다. 오래걸리고 돌아가는 것 같아도 역시 정공법이 짱인가 싶다.

Reference

1. https://catalog.ngc.nvidia.com/orgs/nvidia/containers/tensorrt/tags