End2End ONNX Export Implementation for YOLOv9

[levipereira]

#130
YOLOv9 ONNX End2End with TensorRT Efficient NMS plugin

• Support for End-to-End ONNX Export: Added support for end-to-end ONNX export in export.py and models/experimental.py.

• Model Compatibility: This functionality currently works with all DetectionModel models ;

• Configuration Variables: Use the following flags to configure the model:

  • --include onnx_end2end: Enabled export End2End
  • --topk-all: ONNX END2END/TF.js NMS: Top-k for all classes to keep (default: 100).
  • --iou-thres: ONNX END2END/TF.js NMS: IoU threshold (default: 0.45).
  • --conf-thres: ONNX END2END/TF.js NMS: Confidence threshold (default: 0.25).

Example:

$ python3 export.py --weights ./yolov9-c.pt --include onnx_end2end

export: data=data/coco.yaml, weights=['./yolov9-c.pt'], imgsz=[640], batch_size=1, device=cpu, half=False, inplace=True, keras=False, optimize=False, int8=False, dynamic=True, simplify=True, opset=12, verbose=False, workspace=4, nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, conf_thres=0.25, include=['onnx_end2end']
YOLOv5 🚀 v0.1-27-g86b0667 Python-3.8.10 torch-1.14.0a0+44dac51 CPU

Fusing layers...
Model summary: 604 layers, 50880768 parameters, 0 gradients, 237.6 GFLOPs

PyTorch: starting from ./yolov9-c.pt with output shape (1, 84, 8400) (98.4 MB)

ONNX END2END: starting export with onnx 1.13.0...
/yolov9/models/experimental.py:102: FutureWarning: 'torch.onnx._patch_torch._graph_op' is deprecated in version 1.13 and will be removed in 1.14. Please note 'g.op()' is to be removed from torch.Graph. Please open a GitHub issue if you need this functionality..
  out = g.op("TRT::EfficientNMS_TRT",
[W shape_type_inference.cpp:1913] Warning: The shape inference of TRT::EfficientNMS_TRT type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. (function UpdateReliable)
[W shape_type_inference.cpp:1913] Warning: The shape inference of TRT::EfficientNMS_TRT type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. (function UpdateReliable)
[W shape_type_inference.cpp:1913] Warning: The shape inference of TRT::EfficientNMS_TRT type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. (function UpdateReliable)
[W shape_type_inference.cpp:1913] Warning: The shape inference of TRT::EfficientNMS_TRT type is missing, so it may result in wrong shape inference for the exported graph. Please consider adding it in symbolic function. (function UpdateReliable)
========== Diagnostic Run torch.onnx.export version 1.14.0a0+44dac51 ===========
verbose: False, log level: Level.ERROR
======================= 0 NONE 0 NOTE 4 WARNING 0 ERROR ========================
4 WARNING were not printed due to the log level.


Starting to simplify ONNX...
ONNX export success, saved as ./yolov9-c_end2end.onnx
ONNX END2END: export success ✅ 11.5s, saved as ./yolov9-c_end2end.onnx (129.3 MB)

Export complete (13.6s)
Results saved to /yolov9/experiments/models
Visualize:       https://netron.app

#130 (comment) #188 (comment)

[mdciri]

@WongKinYiu, would you please add to models/experimental.py also the class ORT_NMS and ONNX_ORT as in yolov7 with the right changes? Thanks in advance.