New PR for "ultralytics#7736"

Author: triple-Mu

examples/export.sh

@@ -0,0 +1,16 @@
+#!/usr/bin/env bash
+cd ../
+mkdir -p weights
+
+# download official weights
+wget https://gh.ddlc.top/https://github.com/ultralytics/yolov5/releases/download/v6.1/yolov5s.pt -P weights
+# export yolov5s.onnx
+python3 export.py --weights weights/yolov5s.pt --include onnx engine  --nms
+mv weights/yolov5s.onnx ./examples/yolov5s_nms.onnx
+cd examples
+trtexec --onnx=./yolov5s_nms.onnx --saveEngine=./yolov5s_nms_fp16.engine --fp16
+
+# result test
+wget https://oneflow-static.oss-cn-beijing.aliyuncs.com/tripleMu/image1.jpg
+python3 trt_infer.py
+trtexec --loadEngine=./yolov5s_nms_fp16.engine --verbose --useCudaGraph --noDataTransfers --shapes=images:1x3x640x640

examples/trt_infer.py

@@ -0,0 +1,99 @@
+import sys
+
+import cv2
+
+sys.path.append('../')
+import random
+import time
+from collections import OrderedDict, namedtuple
+
+import numpy as np
+import tensorrt as trt
+import torch
+from PIL import Image
+
+from utils.augmentations import letterbox
+
+names = [
+    'person', 'bicycle', 'car', 'motorcycle', 'airplane', 'bus', 'train', 'truck', 'boat', 'traffic light',
+    'fire hydrant', 'stop sign', 'parking meter', 'bench', 'bird', 'cat', 'dog', 'horse', 'sheep', 'cow', 'elephant',
+    'bear', 'zebra', 'giraffe', 'backpack', 'umbrella', 'handbag', 'tie', 'suitcase', 'frisbee', 'skis', 'snowboard',
+    'sports ball', 'kite', 'baseball bat', 'baseball glove', 'skateboard', 'surfboard', 'tennis racket', 'bottle',
+    'wine glass', 'cup', 'fork', 'knife', 'spoon', 'bowl', 'banana', 'apple', 'sandwich', 'orange', 'broccoli',
+    'carrot', 'hot dog', 'pizza', 'donut', 'cake', 'chair', 'couch', 'potted plant', 'bed', 'dining table', 'toilet',
+    'tv', 'laptop', 'mouse', 'remote', 'keyboard', 'cell phone', 'microwave', 'oven', 'toaster', 'sink', 'refrigerator',
+    'book', 'clock', 'vase', 'scissors', 'teddy bear', 'hair drier', 'toothbrush']
+colors = {name: [random.randint(0, 255) for _ in range(3)] for i, name in enumerate(names)}
+
+w = './yolov5s_nms_fp16.engine'
+image_path = './image1.jpg'
+device = torch.device('cuda:0')
+
+# Infer TensorRT Engine
+Binding = namedtuple('Binding', ('name', 'dtype', 'shape', 'data', 'ptr'))
+logger = trt.Logger(trt.Logger.INFO)
+trt.init_libnvinfer_plugins(logger, namespace="")
+with open(w, 'rb') as f, trt.Runtime(logger) as runtime:
+    model = runtime.deserialize_cuda_engine(f.read())
+bindings = OrderedDict()
+fp16 = False  # default updated below
+for index in range(model.num_bindings):
+    name = model.get_binding_name(index)
+    dtype = trt.nptype(model.get_binding_dtype(index))
+    shape = tuple(model.get_binding_shape(index))
+    data = torch.from_numpy(np.empty(shape, dtype=np.dtype(dtype))).to(device)
+    bindings[name] = Binding(name, dtype, shape, data, int(data.data_ptr()))
+    if model.binding_is_input(index) and dtype == np.float16:
+        fp16 = True
+binding_addrs = OrderedDict((n, d.ptr) for n, d in bindings.items())
+context = model.create_execution_context()
+
+image = cv2.imread(image_path)
+image, ratio, dwdh = letterbox(image, auto=False)
+image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+image_copy = image.copy()
+
+image = image.transpose((2, 0, 1))
+image = np.expand_dims(image, 0)
+image = np.ascontiguousarray(image)
+im = torch.from_numpy(image).to(device)
+im = im.float()
+im /= 255
+
+# warmup for 10 times
+for _ in range(10):
+    tmp = torch.randn(1, 3, 640, 640).to(device)
+    binding_addrs['images'] = int(tmp.data_ptr())
+    context.execute_v2(list(binding_addrs.values()))
+
+start = time.perf_counter()
+binding_addrs['images'] = int(im.data_ptr())
+context.execute_v2(list(binding_addrs.values()))
+print(f'Cost {time.perf_counter()-start} s')
+
+nums = bindings['num_dets'].data
+boxes = bindings['det_boxes'].data
+scores = bindings['det_scores'].data
+classes = bindings['det_classes'].data
+
+print(nums)
+print(boxes)
+print(scores)
+print(classes)
+
+num = int(nums[0][0])
+box_img = boxes[0, :num].round().int()
+score_img = scores[0, :num]
+clss_img = classes[0, :num]
+for i, (box, score, clss) in enumerate(zip(box_img, score_img, clss_img)):
+    name = names[clss]
+    color = colors[name]
+    cv2.rectangle(image_copy, box[:2].tolist(), box[2:].tolist(), color, 2)
+    cv2.putText(image_copy,
+                name, (int(box[0]), int(box[1]) - 2),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.75, [225, 255, 255],
+                thickness=2)
+
+Image.fromarray(image_copy).show()

export.py

@@ -91,7 +91,7 @@ def export_formats():
         ['TensorFlow Lite', 'tflite', '.tflite', True, False],
         ['TensorFlow Edge TPU', 'edgetpu', '_edgetpu.tflite', False, False],
         ['TensorFlow.js', 'tfjs', '_web_model', False, False],
-        ['PaddlePaddle', 'paddle', '_paddle_model', True, True],]
+        ['PaddlePaddle', 'paddle', '_paddle_model', True, True], ]
     return pd.DataFrame(x, columns=['Format', 'Argument', 'Suffix', 'CPU', 'GPU'])
 
 
@@ -185,6 +185,66 @@ def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr('ONNX
     return f, model_onnx
 
 
+@try_export
+def export_onnx_for_backend(model, im, file, opset, nms_cfg, dynamic, simplify, prefix=colorstr('ONNX:')):
+    # YOLOv5 ONNX export
+    check_requirements(('onnx',))
+    import onnx
+
+    LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...')
+    f = file.with_suffix('.onnx')
+
+    from models.common import End2End
+    model = End2End(model, *nms_cfg, device=im.device)
+    if nms_cfg[-1] == 'ort':
+        output_names = ['outputs']
+    elif nms_cfg[-1] == 'trt':
+        output_names = ['num_dets', 'det_boxes', 'det_scores', 'det_classes']
+
+    if dynamic and nms_cfg[-1] == 'ort':
+        dynamic_cfg = {n: {0: 'batch'} for n in output_names}
+    elif dynamic and nms_cfg[-1] == 'trt':
+        dynamic_cfg = {n: {0: 'batch'} for n in output_names}
+
+    torch.onnx.export(
+        model.cpu() if dynamic else model,  # --dynamic only compatible with cpu
+        im.cpu() if dynamic else im,
+        f,
+        verbose=False,
+        opset_version=opset,
+        training=torch.onnx.TrainingMode.EVAL,
+        do_constant_folding=True,
+        input_names=['images'],
+        output_names=output_names,
+        dynamic_axes=dynamic_cfg if dynamic else None)
+
+    # Checks
+    model_onnx = onnx.load(f)  # load onnx model
+    onnx.checker.check_model(model_onnx)  # check onnx model
+
+    # Metadata
+    d = {'stride': int(max(model.stride)), 'names': model.names}
+    for k, v in d.items():
+        meta = model_onnx.metadata_props.add()
+        meta.key, meta.value = k, str(v)
+    onnx.save(model_onnx, f)
+
+    # Simplify
+    if simplify:
+        try:
+            cuda = torch.cuda.is_available()
+            check_requirements(('onnxruntime-gpu' if cuda else 'onnxruntime', 'onnx-simplifier>=0.4.1'))
+            import onnxsim
+
+            LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...')
+            model_onnx, check = onnxsim.simplify(model_onnx)
+            assert check, 'assert check failed'
+            onnx.save(model_onnx, f)
+        except Exception as e:
+            LOGGER.info(f'{prefix} simplifier failure: {e}')
+    return f, model_onnx
+
+
 @try_export
 def export_openvino(file, metadata, half, prefix=colorstr('OpenVINO:')):
     # YOLOv5 OpenVINO export
@@ -447,9 +507,9 @@ def export_tfjs(file, prefix=colorstr('TensorFlow.js:')):
             r'"Identity.?.?": {"name": "Identity.?.?"}, '
             r'"Identity.?.?": {"name": "Identity.?.?"}, '
             r'"Identity.?.?": {"name": "Identity.?.?"}}}', r'{"outputs": {"Identity": {"name": "Identity"}, '
-            r'"Identity_1": {"name": "Identity_1"}, '
-            r'"Identity_2": {"name": "Identity_2"}, '
-            r'"Identity_3": {"name": "Identity_3"}}}', json)
+                                                           r'"Identity_1": {"name": "Identity_1"}, '
+                                                           r'"Identity_2": {"name": "Identity_2"}, '
+                                                           r'"Identity_3": {"name": "Identity_3"}}}', json)
         j.write(subst)
     return f, None
 
@@ -506,6 +566,7 @@ def run(
         verbose=False,  # TensorRT: verbose log
         workspace=4,  # TensorRT: workspace size (GB)
         nms=False,  # TF: add NMS to model
+        backend='ort',  # Backend for export NMS
         agnostic_nms=False,  # TF: add agnostic NMS to model
         topk_per_class=100,  # TF.js NMS: topk per class to keep
         topk_all=100,  # TF.js NMS: topk for all classes to keep
@@ -518,6 +579,7 @@ def run(
     flags = [x in include for x in fmts]
     assert sum(flags) == len(include), f'ERROR: Invalid --include {include}, valid --include arguments are {fmts}'
     jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags  # export booleans
+    end2end, onnx = onnx and nms, onnx and not nms
     file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights)  # PyTorch weights
 
     # Load PyTorch model
@@ -554,7 +616,7 @@ def run(
     LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)")
 
     # Exports
-    f = [''] * len(fmts)  # exported filenames
+    f = [''] * (len(fmts)) + 1  # exported filenames
     warnings.filterwarnings(action='ignore', category=torch.jit.TracerWarning)  # suppress TracerWarning
     if jit:  # TorchScript
         f[0], _ = export_torchscript(model, im, file, optimize)
@@ -592,6 +654,9 @@ def run(
     if paddle:  # PaddlePaddle
         f[10], _ = export_paddle(model, im, file, metadata)
 
+    if end2end:
+        nms_cfg = [topk_all, iou_thres, conf_thres, backend]
+        f[10], _ = export_onnx_for_backend(model, im, file, opset, nms_cfg, dynamic, simplify)
     # Finish
     f = [str(x) for x in f if x]  # filter out '' and None
     if any(f):
@@ -628,6 +693,7 @@ def parse_opt():
     parser.add_argument('--verbose', action='store_true', help='TensorRT: verbose log')
     parser.add_argument('--workspace', type=int, default=4, help='TensorRT: workspace size (GB)')
     parser.add_argument('--nms', action='store_true', help='TF: add NMS to model')
+    parser.add_argument('--backend', type=str, default='ort', help='Backend for export NMS')
     parser.add_argument('--agnostic-nms', action='store_true', help='TF: add agnostic NMS to model')
     parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep')
     parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep')