End2end (#61)

* export end2end onnx model * fixbug * add web demo (#58) * Update README.md * main code update yolov7-tiny deploy cfg * main code update yolov7-tiny training cfg * main code @liguagua752109150 https://github.com/WongKinYiu/yolov7/issues/33#issuecomment-1178669212 * main code @albertfaromatics https://github.com/WongKinYiu/yolov7/issues/35#issuecomment-1178800685 * main code update link * main code add custom hyp * main code update default activation function * main code update path * main figure add more tasks * main code update readme * main code update reparameterization * Update README.md * main code update readme * main code update aux training * main code update aux training * main code update aux training * main figure update yolov7 prediction * main code update readme * main code rename * main code rename * main code rename * main code rename * main code update readme * main code update visualization * main code fix gain for train_aux * main code update loss * main code update instance segmentation demo * main code update keypoint detection demo * main code update pose demo * main code update pose * main code update pose * main code update pose * main code update pose * main code update trace * Update README.md * main code fix ciou * main code fix nan of aux training https://github.com/WongKinYiu/yolov7/issues/250#issue-1312356380 @hudingding * support onnx to tensorrt convert (#114) * fuse IDetect (#148) * Fixes #199 (#203) * minor fix * resolve conflict * resolve conflict * resolve conflict * resolve conflict * resolve conflict * resolve * resolve * resolve * resolve Co-authored-by: AK391 <81195143+AK391@users.noreply.github.com> Co-authored-by: Alexey <AlexeyAB@users.noreply.github.com> Co-authored-by: Kin-Yiu, Wong <102582011@cc.ncu.edu.tw> Co-authored-by: linghu8812 <36389436+linghu8812@users.noreply.github.com> Co-authored-by: Alexander <84590713+SashaAlderson@users.noreply.github.com> Co-authored-by: Ben Raymond <ben@theraymonds.org> Co-authored-by: AlexeyAB84 <alexeyab84@gmail.com>
2022-07-22 21:24:13 +08:00
parent 2596994f39
commit afdc86f519
12 changed files with 283 additions and 17 deletions
@@ -0,0 +1,151 @@
+import numpy as np
+import onnx
+from onnx import shape_inference
+import onnx_graphsurgeon as gs
+import logging
+
+LOGGER = logging.getLogger(__name__)
+
+class RegisterNMS(object):
+    def __init__(
+        self,
+        onnx_model_path: str,
+        precision: str = "fp32",
+    ):
+
+        self.graph = gs.import_onnx(onnx.load(onnx_model_path))
+        assert self.graph
+        LOGGER.info("ONNX graph created successfully")
+        # Fold constants via ONNX-GS that PyTorch2ONNX may have missed
+        self.graph.fold_constants()
+        self.precision = precision
+        self.batch_size = 1
+    def infer(self):
+        """
+        Sanitize the graph by cleaning any unconnected nodes, do a topological resort,
+        and fold constant inputs values. When possible, run shape inference on the
+        ONNX graph to determine tensor shapes.
+        """
+        for _ in range(3):
+            count_before = len(self.graph.nodes)
+
+            self.graph.cleanup().toposort()
+            try:
+                for node in self.graph.nodes:
+                    for o in node.outputs:
+                        o.shape = None
+                model = gs.export_onnx(self.graph)
+                model = shape_inference.infer_shapes(model)
+                self.graph = gs.import_onnx(model)
+            except Exception as e:
+                LOGGER.info(f"Shape inference could not be performed at this time:\n{e}")
+            try:
+                self.graph.fold_constants(fold_shapes=True)
+            except TypeError as e:
+                LOGGER.error(
+                    "This version of ONNX GraphSurgeon does not support folding shapes, "
+                    f"please upgrade your onnx_graphsurgeon module. Error:\n{e}"
+                )
+                raise
+
+            count_after = len(self.graph.nodes)
+            if count_before == count_after:
+                # No new folding occurred in this iteration, so we can stop for now.
+                break
+
+    def save(self, output_path):
+        """
+        Save the ONNX model to the given location.
+        Args:
+            output_path: Path pointing to the location where to write
+                out the updated ONNX model.
+        """
+        self.graph.cleanup().toposort()
+        model = gs.export_onnx(self.graph)
+        onnx.save(model, output_path)
+        LOGGER.info(f"Saved ONNX model to {output_path}")
+
+    def register_nms(
+        self,
+        *,
+        score_thresh: float = 0.25,
+        nms_thresh: float = 0.45,
+        detections_per_img: int = 100,
+    ):
+        """
+        Register the ``EfficientNMS_TRT`` plugin node.
+        NMS expects these shapes for its input tensors:
+            - box_net: [batch_size, number_boxes, 4]
+            - class_net: [batch_size, number_boxes, number_labels]
+        Args:
+            score_thresh (float): The scalar threshold for score (low scoring boxes are removed).
+            nms_thresh (float): The scalar threshold for IOU (new boxes that have high IOU
+                overlap with previously selected boxes are removed).
+            detections_per_img (int): Number of best detections to keep after NMS.
+        """
+
+        self.infer()
+        # Find the concat node at the end of the network
+        op_inputs = self.graph.outputs
+        op = "EfficientNMS_TRT"
+        attrs = {
+            "plugin_version": "1",
+            "background_class": -1,  # no background class
+            "max_output_boxes": detections_per_img,
+            "score_threshold": score_thresh,
+            "iou_threshold": nms_thresh,
+            "score_activation": False,
+            "box_coding": 0,
+        }
+
+        if self.precision == "fp32":
+            dtype_output = np.float32
+        elif self.precision == "fp16":
+            dtype_output = np.float16
+        else:
+            raise NotImplementedError(f"Currently not supports precision: {self.precision}")
+
+        # NMS Outputs
+        output_num_detections = gs.Variable(
+            name="num_detections",
+            dtype=np.int32,
+            shape=[self.batch_size, 1],
+        )  # A scalar indicating the number of valid detections per batch image.
+        output_boxes = gs.Variable(
+            name="detection_boxes",
+            dtype=dtype_output,
+            shape=[self.batch_size, detections_per_img, 4],
+        )
+        output_scores = gs.Variable(
+            name="detection_scores",
+            dtype=dtype_output,
+            shape=[self.batch_size, detections_per_img],
+        )
+        output_labels = gs.Variable(
+            name="detection_classes",
+            dtype=np.int32,
+            shape=[self.batch_size, detections_per_img],
+        )
+
+        op_outputs = [output_num_detections, output_boxes, output_scores, output_labels]
+
+        # Create the NMS Plugin node with the selected inputs. The outputs of the node will also
+        # become the final outputs of the graph.
+        self.graph.layer(op=op, name="batched_nms", inputs=op_inputs, outputs=op_outputs, attrs=attrs)
+        LOGGER.info(f"Created NMS plugin '{op}' with attributes: {attrs}")
+
+        self.graph.outputs = op_outputs
+
+        self.infer()
+
+    def save(self, output_path):
+        """
+        Save the ONNX model to the given location.
+        Args:
+            output_path: Path pointing to the location where to write
+                out the updated ONNX model.
+        """
+        self.graph.cleanup().toposort()
+        model = gs.export_onnx(self.graph)
+        onnx.save(model, output_path)
+        LOGGER.info(f"Saved ONNX model to {output_path}")