[YOLOv8] YOLOv8 Export, Pytorch to TensorRT

이전 포스팅

2023.07.21 - [DeepLearning/YOLO] - [YOLOv8] YOLOv8 install windows 10

[YOLOv8] YOLOv8 install windows 10

---

YOLOv8 Export

- ultralytics에서는 ONNX, OpenVINO, TensorRT 등 다양한 format으로 export가 가능하다.

- 학습된 모델을 최적화하여 NVIDIA GPU에서 추론 속도를 향상시키는 TensorRT로의 export를 소개하려 한다.

 

 

1. TensorRT Install

https://developer.nvidia.com/tensorrt

- Anaconda prompt에서 이전 YOLOv8를 설치한 환경을 활성화. 해당 환경에서 TensorRT를 설치한다.

- 본인은 TensorRT 8.5.1.7을 사용.

 

 

cd graphsurgeon
pip install graphsurgeon-0.4.6-py2.py3-none-any.whl

cd onnx_graphsurgeon
pip install onnx_graphsurgeon-0.3.12-py2.py3-none-any.whl

cd python
pip install tensorrt-8.5.1.7-cp39-none-win_amd64.whl

cd uff
pip install uff-0.6.9-py2.py3-none-any.whl

- yolov8 환경에서 위 명령어를 통해 TensorRT 설치.

- tensorrt-8.5.1.7-cp39-none-win_amd64.whl의 경우 자신의 Python version에 맞게 설치해야 한다.

 

- tensorrt 설치 확인.

 

 

2. Export TensorRT

from ultralytics import YOLO

# Load a model
model = YOLO('yolov8n.pt')

# Export the model
model.export(format='engine', device=0)

 

- format='engine'으로 pytorch model을 TensorRT로 export.

 

 

3. Python 환경에서 Predict Test

 

from ultralytics import YOLO
import cv2


# Load a model
model = YOLO('yolov8n.engine')  # pretrained YOLOv8n model

# Load a image
image = cv2.imread('./test_image.jpg')

# Run inference
results = model(image, conf=0.4)

# Draw results
result_plotted = results[0].plot()

cv2.imshow("result", result_plotted)
cv2.waitKey(0)

- 위에서 생성한 .engine model을 불러와 predict.

 

Pytorch vs TensorRT

 

Pytorch (.pt) 결과

TensorRT (.engine) 결과

 

- 두 모델을 같은 이미지로 실행한 결과 inference time이 13.0ms -> 1.0ms로 줄어든 것을 확인할 수 있다.

 

 

 

 

 

 

끝.