Ubuntu, TensorRT로 Yolov5 inference

Ubuntu 20.04

TensorRT 8.4.3.1

 

공식 깃헙:

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

GitHub - ultralytics/yolov5: YOLOv5 🚀 in PyTorch > ONNX > CoreML > TFLite

 

공식 깃헙에 따르면, 

Yolov5를 TensorRT를 통해 배포하는 방법에는

1. export.py를 거쳐 model.engine을 만든 후, detect.py 에 로드하여 돌리기

 

!python detect.py --weights yolov5s.engine --source ../zidane.jpg

 

2. export.py를 거쳐 model.engine을 만든 후, torch.hub.load에 로드해서 돌리기

 

model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.engine', force_reload=True)

 

가 있는데,

각각의 방법과 2번에서 발생하는 input shape 버그(2022.12.05 기준) 해결에 대해서 정리함

TensorRT가 설치되어 있다는 것을 가정 (Ubuntu, TensorRT 설치 (tistory.com))

 

1. detect.py 에 로드하여 돌리기

 

1-1. 공식 깃헙 다운로드

 

# Jupyter에서 돌릴 경우 맨 앞에 ! 추가
git clone https://github.com/ultralytics/yolov5.git

 

1-2. 필수 패키지 설치

 

# 다운로드 한 yolov5 폴더 안에서
python3 -m pip install -r requirements.txt

 

1-3. pytorch 모델을 tensorrt 로 변환

(GPU를 사용하며, 시간이 조금 걸릴 수 있다.)

 

# yolov5s.pt 를 yolov5s.engine 으로 변환 
python export.py --weights yolov5s.pt --include engine --device 0

 

1-4. inference

 

# yolov5s.engine 으로 zidane.jpg 파일에 inference
python detect.py --weights yolov5s.engine --source ../zidane.jpg

 

1-5. 결과 확인

 

# yolov5 폴더 안에서!
result_path = 'runs/detect/exp/zidane.jpg'

from IPython.display import Image
Image(filename=result_path)

 

output:

 

 

2. torch.hub.load를 이용하기

 

2-1. 모델 정의(불러오기)

 

import torch

model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.engine', force_reload=True)

 

2-2. 이미지 경로 정의

 

img_path = 'zidane.jpg'

 

2-3. inference

 

result = model(img_path)

 

원래는 위 코드로 간단히 돌아야 한다.

2022.12.05 기준, tensorrt 모델을 불러와서 처리하는 과정에 버그가 있는지

아래와 같이 "not equal to max model size" 에러가 발생한다.

resize 과정에 문제가 발생한듯 하다.

 

 

해결방법:

model inference에 앞서 resizing을 하고, model에 size 인자를 넣어준다.

 

# default size는 640
image = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)
resized_image = cv2.resize(image, (416, 416))

results = model(resized_image, size=416)

 

끝