JetsonTX2 | Pytorch Model을 TensorRT로 변환, torch to trt, torch to onnx, onnx to engine

JetPack 4.6.2 설치 방법 : https://iambeginnerdeveloper.tistory.com/228

JetPack 4.6.2에서 python3.8 설치하는 방법 : https://iambeginnerdeveloper.tistory.com/229

python3.8용 TensorRT 설치하는 방법 : https://iambeginnerdeveloper.tistory.com/239

---

드디어 jetsonTX2에서 trt 파일로 yolov8을 실행 시키기 위한 마지막 단계이다.

(TensorRT는 NVIDIA에서 만든 inference engine으로, kernel fusion, graph optimization, low precision 등의 optimization을 도와준다.)

JetsonTX2에서 yolov8을 조금 더 빠르게 실행 시키기 위해서 할 수 있는 방법은 pytorch model을 TensorRT로 변환하여 동작 시키는 것이다.

이것을 위해 jetpack을 설치하고, python3.8을 따로 설치 해 주고.. tensorrt까지 새로 설치 한 것이다.

 

 

보통, pytorch 모델을 onnx로 변환을 먼저 한 뒤, onnx를 tensorrt로 변환 해 주는 것이 일반적이다.

 

1. pytorch to onnx

pytorch에서 onnx 변환은 Jetson에서 yolov8의 export.py 코드로 충분히 가능하다.

cd yolov8/

python3.8 export.py

yolov8 디렉토리로 이동해서 export.py를 실행 시켜 주면 된다.

# export.py

from ultralytics import YOLO

# Load a model
model = YOLO('/home/aivis/workspace/yolov8/weights/230426_yolov8n_best.pt')

# Export the model
model.export(format='onnx', imgsz=640, half=True)

경로를 맞춰서 코드를 실행 시켜 주면 onnx 파일이 생성되어 있을 것이다.

 

2. onnx to engine

Jetson에서 onnx 파일을 engine(trt) 파일로 변환하는 가장 간단한 방법은 jetson 자체에 있는 trtexec tools을 사용하는 것이다.

# Using trtexec tools for export engine
/usr/src/tensorrt/bin/trtexec \
--onnx=yolov8s.onnx \
--saveEngine=yolov8s.engine

 위 명령어를 실행 하고 나면 saveEngine에 넣은 경로에 변환된 engine 파일이 저장되어 있을 것이다.

 

이렇게 생성한 engine 파일을 weight 경로에 넣어주고 yolov8을 실행 시키면 문제없이 실행이 될 것이다.

 

---

보통은 위처럼 두 단계를 걸쳐서 tensorrt로 변환하지만, 

https://github.com/pytorch/TensorRT

GitHub - pytorch/TensorRT: PyTorch/TorchScript/FX compiler for NVIDIA GPUs using TensorRT

친절하게도 한번에 pytorch에서 tensorrt로 변환 해 주는 컴파일러가 있다.

👍