TensorRT를 이용해서 Yolov4 성능 향상시키기

Jetson Xavier에서 Yolov4로 실시간 객체 인식을 하면 14프레임 정도나옵니다.

 

여기서 딥러닝 네트워크를 최적화 시켜주는 도구인 TensorRT를 이용하면 프레임이 더 잘 나오게 됩니다.

 

방법은 깃허브를 참고했습니다.

https://github.com/jkjung-avt/tensorrt_demos

GitHub - jkjung-avt/tensorrt_demos: TensorRT MODNet, YOLOv4, YOLOv3, SSD, MTCNN, and GoogLeNet

---

git clone https://github.com/jkjung-avt/tensorrt_demos
cd ~/tensorrt_demos/yolo
./install_pycuda.sh
sudo pip3 install onnx==1.4.1
cd ~/tensorrt_demos/plugins
make
cd ~/tensorrt_demos/yolo
./download_yolo.sh

맨 마지막 명령어는 여러가지 yolo의 cfg파일과 weights파일을 다운 받는 명령어입니다.

 

여기서 제공하는 파일을 쓰셔도 되고 만약 본인이 직접 학습한 파일이 있다면 이 경로에 자신이 사용할 cfg파일와 weights파일을 넣어줍니다. (두 파일의 이름은 확장자명만 다르고 동일해야 합니다!)

 

---

 

그리고 yolo파일을 tensorrt로 바꿔주기 위해 딥러닝 네트워크의 표준 포맷인 onnx로 먼저 변환을 해줍니다.

python3 yolo_to_onnx.py -m yolov4

맨 뒤에 yolov4가 본인이 변경하기 원하는 파일의 이름을 넣어주면 됩니다.

 

 

그리고 onnx를 다시 tensorrt로 바꿔줍니다.

python3 onnx_to_tensorrt.py -m yolov4

 

변환을 다 마치고 나면 

 

yolo디렉토리에 trt파일이 생성된 것을 볼 수 있습니다.

 

이제 trt파일을 실행해봅시다.

 cd ~/tensorrt_demos
 wget https://raw.githubusercontent.com/pjreddie/darknet/master/data/dog.jpg -O ${HOME}/Pictures/dog.jpg
 python3 trt_yolo.py --image ~/Pictures/dog.jpg \ -m yolov4-416
 python3 trt_yolo.py --usb 0 -m yolov4 --width 640 --height 480

세번째 명령어는 사진을 인식하는 것이고

마지막 명령어는 웹캠을 이용하는 명령어입니다.

명령어를 잘 보고 인식하고자 하는 대상의 위치나 본인이 사용하고자하는 trt파일의 이름을 적당히 넣어주면 됩니다.

 

640*480 크기에서 30프레임 정도 안정적으로 잘 나옵니다.

 

만약 본인이 custom한 yolo파일을 사용하는 경우에는

tensorrt_demos/utils에 yolo_classes.py 파일에 들어가 class_list를 변경해주면 됩니다.