[TensorRT] TensorRT Processing using CuPy(with Yolov5)

개요

Yolov5 네트워크를 TensorRT로 최적화하여 객체감지를 돌리는 프로젝트에서 기존 객체감지 모델과 마찬가지로 입력 이미지를 전처리하는 것에 NumPy를, GPU 메모리에 입력과 출력 할당에 PyCUDA 라이브러리를 사용하였다. 이제는 CuPy를 이용하여 GPU기반 연산 및 메모리 사용으로 더 빠르게 수행해보는 내용을 정리한다.

기존 NumPy + PyCuda

집중적으로 살펴볼 곳은 NumPy를 통한 이미지 전처리 부분과 PyCUDA를 사용한 입출력 GPU 메모리 할당 부분이다. Yolov5 객체감지를 기반에 두고 있음을 참고한다.

입출력 할당

GPU 메모리에 입력(inputs)과 출력(outputs)에 할당하는 배열 영역을 생성하고 그 주소를 bindings 에 묶어준다.

• PyCUDA 라이브러리를 통해 GPU 메모리에 기반한 배열 생성

"""
각종 tensorrt 튜토리얼이나 예제 코드에서 allocate_buffer()에 해당하는 부분이다.
"""
import pycuda.driver as cuda

inputs = []
outputs = []
bindings = []

for binding in engine:
    size = trt.volume(engine.get_binding_shape(binding)) * engine.max_batch_size
    dtype = trt.nptype(engine.get_binding_dtype(binding))

    host_mem = cuda.pagelocked_empty(size, dtype)
    device_mem = cuda.mem_alloc(host_mem.nbytes)
    # Append the device buffer to device bindings.
    # int(device_mem) is memory address
    bindings.append(int(device_mem))

    # HostDeviceMem class is simple namespace
    if engine.binding_is_input(binding):
        inputs.append(HostDeviceMem(host_mem, device_mem))
    else:
        outputs.append(HostDeviceMem(host_mem, device_mem))

전처리

요약된 전처리 과정은 다음과 같다.

• 이미지 리사이즈 (비율 유지 한 채로 기준 크기에 맞게)
• 이미지 패딩 (정사각형에 맞게 모자란 곳 보더 생성)
• transpose (HWC to CHW)
• type cast (unit8 to flaot32)
• divide ([0,255] to [0,1])
• sort c-array(메모리에 배열 정렬)
• GPU 메모리(입력)에 이미지 복사

"""
일반적인 cv tensorrt 연산에서 preprocessing에 해당하는 부분이다.
"""
import numpy as np
import cv2

# image resize (tw, th는 yolov5 전처리에 맞는 계산값)
img = cv2.resize(image, (tw, th))

# Pad the short side with (128,128,128)
img = cv2.copyMakeBorder(img, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, None, (128, 128, 128))

# HWC to CHW format
img = np.transpose(img, [2, 0, 1])

# Type Cast
img = img.astype(np.float32)

# Normalize to [0,1]
img /= 255.0

# Convert the image to row-major order, also known as "C order":
img = np.ascontiguousarray(img)

# Image COPY to input memory(GPU Memory)
np.copyto(inputs[0].host, img.ravel())

추론(Inference)

네트워크 모델에 대해 입력 배열을 가지고 추론 시에도 PyCUDA의 일정 개입이 들어간다. 모듈의 이름이 혼동되는 것을 막기 위해 import는 같이 명시한다. htod, dtoh 등의 메모리 복사 연산들이 수행되는 것으로 보아 cpu 메모리(RAM)과 gpu 메모리 간의 이동이 주기적으로 일어날 것 임을 알 수 있다.

"""
각종 tensorrt 튜토리얼이나 예제 코드에서 do_inference()에 해당하는 부분이다.
"""
import pycuda.driver as cuda

cuda.init()
device = cuda.Device(gpu_num)
# PyCUDA create ctx(context)
ctx = device.make_context()
# TensorRT engine create context
context = engine.create_execution_context()

ctx.push()
# Transfer input data to the GPU.
[cuda.memcpy_htod_async(inp.device, inp.host, stream) for inp in inputs]
# Run inference.
context.execute_async(batch_size=batch_size, bindings=bindings, stream_handle=stream.handle)
# Transfer predictions back from the GPU.
[cuda.memcpy_dtoh_async(out.host, out.device, stream) for out in outputs]
# Synchronize the stream
stream.synchronize()
ctx.pop()

# Return only the host outputs.
res = [out.host for out in outputs]

CuPy

CuPy 라이브러리는 NumPy와 SciPy를 대체하기 위해 개발된 라이브러리라 opencv를 대체하기엔 아직 조금 문제가 있다. 또한 PyCUDA도 공식 도큐먼트에 언급은 되어있지 않지만 GPU 기반 메모리 할당과 배열을 사용한다는 점에서 충분히 대체가 가능하다.

입출력 할당

위의 전처리와 비슷하지만 호스트와 디바이스 메모리를 할당하는 부분을 Cupy가 처리한다.

inputs = []
outputs = []
bindings = []

for binding in engine:
    # Append the device buffer to device bindings.
    shape = engine.get_binding_shape(binding)
    dtype = trt.nptype(engine.get_binding_dtype(binding))
    # Allocate buffers
    cuda_mem = cp.zeros(shape=shape, dtype=dtype)
    bindings.append(cuda_mem.data.ptr)
    
    if engine.binding_is_input(binding):
        inputs.append(cuda_mem)
    else:
        outputs.append(cuda_mem)

전처리

NumPy 과정과 전처리는 전반적으로 동일하다. 배열을 다루기 전의 CuPy의 ndarray 타입으로 변경 후 NumPy 대신 CuPy 라이브러리로 연산들을 수행한다. 마지막 copyto 수행 시 img.ravel()이 없는 것을 볼 수 있는데, 기존 PyCUDA로 바인딩에 추가한 입력(input)은 binding의 shape을 trt.volume으로 감싸서 1차원의 직렬 바인딩으로 메모리영역을 할당하는 반면, CuPy의 전처리는 입력의 shape 자체로 cp.zeros 배열을 바인딩에 추가하기에 3차원 혹은 배치를 포함한 4차원의 shape로 메모리영역을 할당한다.

import cupy as cp
import cv2

# image resize (tw, th는 yolov5 전처리에 맞는 계산값)
img = cv2.resize(image, (tw, th))

# Pad the short side with (128,128,128)
img = cv2.copyMakeBorder(img, ty1, ty2, tx1, tx2, cv2.BORDER_CONSTANT, None, (128, 128, 128))

# cpu array to gpu array
img = cp.asarray(img)

# HWC to CHW format
img = cp.transpose(img, [2, 0, 1])

# Type Cast
img = img.astype(cp.float32)

# Normalize to [0,1]
img /= 255.0

# Convert the image to row-major order, also known as "C order":
img = cp.ascontiguousarray(img)

# adjust shape for input shape (batch)
img = cp.expand_dims(img, axis=0)

# Image COPY to binding from cupy array
cp.copyto(self.inputs[0], img)

추론(Inference)

기존 PyCUDA 추론 방식에 비해 상당히 간단하고 코드도 몇 줄 안된다.

"""
각종 tensorrt 튜토리얼이나 예제 코드에서 do_inference()에 해당하는 부분이다.
"""
import pycuda.driver as cuda

with cp.cuda.Stream(non_blocking=False) as stream:
    # Run inference.
    self.context.execute_async(bindings=self.bindings, stream_handle=stream.ptr)

    # Synchronize the stream
    stream.synchronize()

    # Return the outputs.
    res = outputs[0]

입력 및 출력에 대한 shape 비교

기존 PyCUDA는 입력과 출력 배열에 대한 메모리 할당을 1차원으로 수행했다. 예를 들어 배치 사이즈가 1이고 입력 형상이 (3, 1280, 1280) 인 Yolov5 네트워크는 1 * 3 * 1280 * 1280을 수행하여 49152000 크기의 영역을 GPU 메모리에 할당한다. 또한, 아직 NMS를 수행하지 않은 원시 출력의 형상 (1, 102000, 10)도 1 * 102000 * 10으로 1020000 크기의 영역을 GPU 메모리에 할당한다. 이를 사용자가 추론의 전과 후에 입력과 출력을 전처리, 후처리하여 사용한다.

그에 비해, CuPy는 TensorRT 엔진이 요구하는 입력과 출력의 형상을 그대로 가져간다. 위와 같은 엔진이라면 입력에 대해 (1, 3, 1280, 1280)의 CuPy - ndarray를 만들어 GPU에 할당하고, 출력에 대해 (1, 102000, 10)의 배열을 만들어 GPU에 할당한다.

이렇기 때문에, PyCUDA와 CuPy의 전처리, 후처리는 조금 다를 수 밖에 없음에 주의한다.

참고: CuPy 개발자 문서
참고: CuPy github issue