EdgeYOLO: anchor-free, edge-friendly

Note

This is a trial version without training code and evaluate code. we will publish complete source code after we publish our first edgeyolo paper.

Intro

• In embeded device such as Nvidia Jetson AGX Xavier, EdgeYOLO reaches 34FPS with 50.6% AP in COCO2017 dataset and 25.9% AP in VisDrone2019 (image input size is 640x640, batch=16, post-process included). And for smaller model EdgeYOLO-S, it reaches 53FPS with 44.1% AP and 63.3% AP^0.5(SOTA in P5 small models) in COCO2017.
• we provide a more effective data augmentation during training.
• small object and medium object detect performace is imporved by using RH loss during the last few training epochs.

Models & Benchmark

• models trained on COCO2017-train

some errors occured in our GPU server, we are trying to fix it and we have to delay the release of EdgeYOLO-M model.

Model	Size	mAPval 0.5:0.95	mAPval 0.5	FPSAGX Xavier trt fp16 batch=16 include NMS	Params train / infer (M)	Download
EdgeYOLO-Tiny-LRELU	416 640	33.1 37.8	50.5 56.7	206 109	7.6 / 7.0	github
EdgeYOLO-Tiny	416 640	37.2 41.4	55.4 60.4	136 67	5.8 / 5.5	github
EdgeYOLO-S	640	44.1	63.3	53	9.9 / 9.3	github
EdgeYOLO-M (ETA: ? days)	640	-	-	-	19.0 / 17.8	-
EdgeYOLO	640	50.6	69.8	34	41.2 / 40.5	github

• models trained on VisDrone2019 (pretrained on COCO2017-train)

Model	Size	mAPval 0.5:0.95	mAPval 0.5	Download
EdgeYOLO-Tiny-LRELU	416 640	12.1 18.5	22.8 33.6	github
EdgeYOLO-Tiny	416 640	14.9 21.8	27.3 38.5	github
EdgeYOLO-S	640	23.6	40.8	github
EdgeYOLO-M (coming soon)	640	-	-	-
EdgeYOLO	640	25.9	43.9	github

Some of our detect results in COCO2017

Coming Soon

• train code. After our paper is released, we will publish train code.
• evaluate code. After our paper is released, we will publish evaluate code.

Quick Start

setup

git clone https://github.com/LSH9832/edgeyolo.git
cd edgeyolo
pip install -r requirements.txt

if you use tensorrt, please make sure torch2trt is installed

git clone https://github.com/NVIDIA-AI-IOT/torch2trt.git
cd torch2trt
python setup.py install

inference

First download weights here

python detect.py --weights edgeyolo_coco.pth --source XXX.mp4 --fp16

# full commands
python detect.py --weights edgeyolo_coco.pth 
                 --source /XX/XXX.mp4     # or dir with images, such as /dataset/coco2017/val2017    (jpg/jpeg, png, bmp, webp is available)
                 --conf-thres 0.25 
                 --nms-thres 0.5 
                 --input-size 640 640 
                 --batch-size 1 
                 --save-dir ./img/coco    # if you press "s", the current frame will be saved in this dir
                 --fp16 
                 --no-fuse                # do not fuse layers
                 --no-label               # do not draw label with class name and confidence

It is recomended to use batch_detect.py with the same commands if batch size > 1

python batch_detect.py --weights edgeyolo_coco.pth --source XXX.mp4 --batch-size 16 --fp16
                       --fps 30    # max fps limitation(new function)

export onnx & tensorrt

python export_pth2onnx.py --weights edgeyolo_coco.pth --simplify

# full commands
python export_pth2onnx.py --weights edgeyolo_coco.pth 
                          --img-size 640 640 
                          --batch 1
                          --opset 11
                          --simplify

it generates file yolo_export/onnx/edgeyolo_coco_640x640_batch1.onnx and yolo_export/onnx/edgeyolo_coco_640x640_batch1.yaml

# (workspace: GB)
python export_onnx2trt.py --onnx yolo_export/onnx/edgeyolo_coco_640x640_batch1.onnx 
                          --yaml yolo_export/onnx/edgeyolo_coco_640x640_batch1.yaml 
                          --workspace 10 
                          --fp16

it will generate

yolo_export/tensorrt/edgeyolo_coco_640x640_batch1.pt         # for python inference
yolo_export/tensorrt/edgeyolo_coco_640x640_batch1.engine     # for c++ inference
yolo_export/tensorrt/edgeyolo_coco_640x640_batch1.txt        # for c++ inference

for python inference

python detect.py --trt --weights yolo_export/tensorrt/edgeyolo_coco_640x640.pt --source XXX.mp4

# full commands
python detect.py --trt 
                 --weights yolo_export/tensorrt/edgeyolo_coco_640x640_batch1.pt 
                 --source XXX.mp4
                 --legacy         # if "img = img / 255" when you train your train model
                 --use-decoder    # if use original yolox tensorrt model before version 0.3.0

It is also recomended to use batch_detect.py with the same commands if batch size > 1

python batch_detect.py --trt --weights yolo_export/tensorrt/edgeyolo_coco_640x640_batch1.pt --source XXX.mp4 --fp16
                       --fps 30    # max fps limitation(new function)

for c++ inference

it will be comming soon