123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155 |
- import numpy as np
- import onnx
- from onnx import shape_inference
- try:
- import onnx_graphsurgeon as gs
- except Exception as e:
- print('Import onnx_graphsurgeon failure: %s' % e)
- 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_dets",
- dtype=np.int32,
- shape=[self.batch_size, 1],
- ) # A scalar indicating the number of valid detections per batch image.
- output_boxes = gs.Variable(
- name="det_boxes",
- dtype=dtype_output,
- shape=[self.batch_size, detections_per_img, 4],
- )
- output_scores = gs.Variable(
- name="det_scores",
- dtype=dtype_output,
- shape=[self.batch_size, detections_per_img],
- )
- output_labels = gs.Variable(
- name="det_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}")
|