MCPcopy Create free account
hub / github.com/NVIDIA/TensorRT / sanitize

Method sanitize

samples/python/detectron2/create_onnx.py:114–142  ·  view source on GitHub ↗

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.

(self)

Source from the content-addressed store, hash-verified

112 self.batch_size = None
113
114 def sanitize(self):
115 """
116 Sanitize the graph by cleaning any unconnected nodes, do a topological resort, and fold constant inputs values.
117 When possible, run shape inference on the ONNX graph to determine tensor shapes.
118 """
119
120 for i in range(3):
121 count_before = len(self.graph.nodes)
122 self.graph.cleanup().toposort()
123 try:
124 for node in self.graph.nodes:
125 for o in node.outputs:
126 o.shape = None
127 model = gs.export_onnx(self.graph)
128 model = shape_inference.infer_shapes(model)
129 self.graph = gs.import_onnx(model)
130 except Exception as e:
131 log.info("Shape inference could not be performed at this time:\n{}".format(e))
132 try:
133 self.graph.fold_constants(fold_shapes=True)
134 except TypeError as e:
135 log.error("This version of ONNX GraphSurgeon does not support folding shapes, please upgrade your "
136 "onnx_graphsurgeon module. Error:\n{}".format(e))
137 raise
138
139 count_after = len(self.graph.nodes)
140 if count_before == count_after:
141 # No new folding occurred in this iteration, so we can stop for now.
142 break
143
144 def get_anchors(self, sample_image):
145 """

Callers 2

update_preprocessorMethod · 0.95
process_graphMethod · 0.95

Calls 6

toposortMethod · 0.80
cleanupMethod · 0.45
infer_shapesMethod · 0.45
infoMethod · 0.45
fold_constantsMethod · 0.45
errorMethod · 0.45

Tested by

no test coverage detected