Darknet(yolo)모델을 TensorRT로 변환하기

그래픽카드가 크게 발전하면서 인공지능은 탄탄대로를 걷고 있다.

심지어 NVIDIA에서 Xavier계열의 보드를 사용할 때, TensorRT 모델을 지원해주기 때문에

임베디드에도 인공지능 모델을 쉽게 최적화하여 넣을 수 있다.

 

이번 글에서는 다크넷 기반의 Yolo모델들을 TensorRT로 변환하는 법에 대해서 공부해보고자 한다.

**참고로 필자는 NVIDIA AGX Xavier 환경에서 실행하였다.

---

사실 말이 거창하지 이미 Converting Tool은 능력자님들께서 github에 올려놓으셨다. 

이 글은 다음 github를 참고하면서 만들었다.

github.com/jkjung-avt/tensorrt_demos

jkjung-avt/tensorrt_demos

YOLO 모델 다운로드

먼저 위의 github를 clone하여주자.

clone을 다했다면 tensorrt_demos라는 폴더가 생겼을 것이다.

 

yolo폴더로 이동하자!

aistudy@aistudy-pc:~/tensorrt_demos/yolo$ 

yolo 폴더로 이동하면 download_yolo.sh파일이 있을 텐데, 이 파일을 실행해주면 yolov3, yolov4, tiny모델 등을 다운로드 해온다.

**만약 본인의 cfg파일이 있다면 다음으로 넘어가면 된다.

 

환경 세팅

환경은 친절하게도 yolo폴더 안에 requirements.txt파일로 넣어놓으셨다.

 

onnx==1.4.1로 나오는데 만약 개인적으로 설치하여

1.8.9등의 높은 버전을 설치했을 시, onnx converting에서 에러가 나올 수 있으므로,

1.4.1버전을 사용하는 것을 추천한다.

 

모델 컨버팅

이제 Yolo를 onnx모델로 바꾸어주는 일만 남았다. 

yolo폴더 내의 yolo_to_onnx.py와 onnx_to_trt.py 두 파이썬 파일을 사용할 것이다.

 

먼저 yolo를 onnx 형태로 바꾸어보자.

aistudy@aistudy-pc:~/tensorrt_demos/yolo$ python3 yolo_to_onnx.py -m yolov4-tiny-416 -c 80

-m: 모델명, -c: 클래스 개수

 

여기서 주의 깊게 봐야 되는 부분이 있다. 바로 모델이름이다. 

무조건 위의 형식대로 cfg이름을 바꾸어주어야한다.

 

(yolov4-tiny) + (416) 이렇게 나누어서 모델 + input shape를 세팅하기 때문이다.

width, height도 서로 다르다면 다음과 같이 입력해주면 된다 => yolov4-tiny-416x544

 

class num도 모델과 맞지 않으면 오류문이 나오니 정확하게 입력하자!

 

**21년도 5월 3일에 올라온 글에 따르면, 뒤에 width height를 붙이지 않아도 자동으로 인식해주는 걸로 바뀐 것 같다. -c옵션도 쓰지 않아도 된다고 한다. 단, weight와 cfg파일의 이름은 같아야한다.

 

위의 명령어를 실행해보면 cfg의 layer관련된 문구가 주르륵 쏟아지면서 마지막에

Checking ONNX model...
Saving ONNX file...
Done.

이런 문구가 나온다면 일단 onnx 파일 변환에는 성공한 것이다!!

 

마지막으로 onnx를 trt모델로 바꾸면 된다!

aistudy@aistudy-pc:~/tensorrt_demos/yolo$ python3 onnx_to_tensorrt.py -m yolov4-tiny-416 -c 80

yolo_to_onnx와 형식이 똑같다. int8형식으로 변환을 원한다면 --int8 옵션을 추가해주면 된다.

**필자가 yolov4-tiny기준으로 int8과 float16 모두 시도해 보았지만 큰 속도 차이가 없었다. 더 큰 모델이면 영향이 클 수도 있겠지만 말이다. 

 

정상적으로 동작했다면

Loading the ONNX file...
Adding yolo_layer plugins...
Building an engine.  This would take a while...
(Use "--verbose" or "-v" to enable verbose logging.)
Completed creating engine.
Serialized the TensorRT engine to file: yolov4-tiny-416.trt

 

이러한 문구들이 나올 것이다. 이럼 trt님께서 폴더에 존재하시는 걸 볼 수 있을 것이다.

 

필자가 테스트해 본 결과 yolov4-tiny기준으로 fps가 약 1.5배~2배정도 차이가 나는 것 같다.