先在 Jetson Nano 上 git clone trt_yolov4-tiny 的 repo (裡面包含我訓練好的模型,共有兩個分別是使用 kaggle 人臉資料集和 WIDIER FACE 資料集 )
$ git clone https://github.com/d246810g2000/trt_yolov4-tiny.git
需要安裝 ONNX 為 1.4.1 的版本,安裝前需要先自行安裝 protobuf-3.8.0 需要大約一小時的時間
$ bash install_protobuf-3.8.0.sh
$ pip3 install onnx==1.4.1
進入 plugins 資料夾裡並建立 yolo_layer 的 plugin,make 完後會得到 libyolo_layer.so 檔,這會用在將 ONNX 模型轉換成 TensorRT 模型上
$ cd ~/trt_yolov4-tiny/plugins
$ make -j4
將訓練好的 yolov4-tiny 人臉偵測模型和 cfg 放入 yolo 資料夾中,由於我訓練的 input_size 是 288 因此改名成 yolov4-tiny-288.weights, yolov4-tiny-288.cfg,數字的部分取決於你訓練的 input_size 為何,如何訓練可以回去看這篇
到 yolo 資料夾中將 yolo 模型轉換成 ONNX 檔,再轉換成 trt 檔
- yolov4-tiny-288.weights, yolov4-tiny-288.cfg 轉換成 yolov4-tiny-288.onnx
- yolov4-tiny-288.onnx 轉換成 yolov4-tiny-288.trt
$ cd ~/trt_yolov4-tiny/yolo
$ python3 yolo_to_onnx.py -c 1 -m yolov4-tiny-288
$ python3 onnx_to_tensorrt.py -c 1 -m yolov4-tiny-288
在 Jetson Nano 上查看 trt_yolo.py 的使用指令
$ python3 trt_yolo.py -h
- 照片的推論
$ python3 trt_yolo.py --image test/test.jpg -c 1 -m yolov4-tiny-288
- 影片的推論
$ python3 trt_yolo.py --video test/test.mp4 --video_looping -c 1 -m yolov4-tiny-288
- camera 的推論
$ python3 trt_yolo.py --onboard 0 -c 1 -m yolov4-tiny-288
由於訓練集比較簡單一點所以推論的效果在偵測小目標上會比較不準,但是可以看到 FPS 高達 40,TensorRT 的加速還是相當有效果的。
我們來看看使用 WIDER FACE 資料集 訓練,training data 共有 12879 張,validation data 共有 3225 張,下面是訓練歷史走向。
照剛剛的步驟,將 weights 和 cfg 轉成 onnx 再轉換成 trt 檔後就會得到 yolov4-tiny-wider-288.trt ,我們來看看效果如何
$ python3 trt_yolo.py --video test/test.mp4 --video_looping -c 1 -m yolov4-tiny-wider-288
可以看到在偵測小物件上還是有很好的效果。