| 35 | |
| 36 | |
| 37 | class EfficientDetGraphSurgeon: |
| 38 | def __init__(self, saved_model_path): |
| 39 | """ |
| 40 | Constructor of the EfficientDet Graph Surgeon object, to do the conversion of an EfficientDet TF saved model |
| 41 | to an ONNX-TensorRT parsable model. |
| 42 | :param saved_model_path: The path pointing to the TensorFlow saved model to load. |
| 43 | """ |
| 44 | saved_model_path = os.path.realpath(saved_model_path) |
| 45 | assert os.path.exists(saved_model_path) |
| 46 | |
| 47 | # Use tf2onnx to convert saved model to an initial ONNX graph. |
| 48 | graph_def, inputs, outputs = tf_loader.from_saved_model( |
| 49 | saved_model_path, None, None, "serve", ["serving_default"] |
| 50 | ) |
| 51 | log.info("Loaded saved model from {}".format(saved_model_path)) |
| 52 | with tf.Graph().as_default() as tf_graph: |
| 53 | tf.import_graph_def(graph_def, name="") |
| 54 | with tf_loader.tf_session(graph=tf_graph): |
| 55 | onnx_graph = tfonnx.process_tf_graph(tf_graph, input_names=inputs, output_names=outputs, opset=11) |
| 56 | onnx_model = optimizer.optimize_graph(onnx_graph).make_model("Converted from {}".format(saved_model_path)) |
| 57 | self.graph = gs.import_onnx(onnx_model) |
| 58 | assert self.graph |
| 59 | log.info("TF2ONNX graph created successfully") |
| 60 | |
| 61 | # Fold constants via ONNX-GS that TF2ONNX may have missed |
| 62 | self.graph.fold_constants() |
| 63 | |
| 64 | # Try to auto-detect by finding if nodes match a specific name pattern expected for either of the APIs. |
| 65 | self.api = None |
| 66 | if len([node for node in self.graph.nodes if "class_net/" in node.name]) > 0: |
| 67 | self.api = "AutoML" |
| 68 | elif len([node for node in self.graph.nodes if "/WeightSharedConvolutionalClassHead/" in node.name]) > 0: |
| 69 | self.api = "TFOD" |
| 70 | assert self.api |
| 71 | log.info("Graph was detected as {}".format(self.api)) |
| 72 | |
| 73 | def sanitize(self): |
| 74 | """ |
| 75 | Sanitize the graph by cleaning any unconnected nodes, do a topological resort, and fold constant inputs values. |
| 76 | When possible, run shape inference on the ONNX graph to determine tensor shapes. |
| 77 | """ |
| 78 | for i in range(3): |
| 79 | count_before = len(self.graph.nodes) |
| 80 | |
| 81 | self.graph.cleanup().toposort() |
| 82 | try: |
| 83 | for node in self.graph.nodes: |
| 84 | for o in node.outputs: |
| 85 | o.shape = None |
| 86 | model = gs.export_onnx(self.graph) |
| 87 | model = shape_inference.infer_shapes(model) |
| 88 | self.graph = gs.import_onnx(model) |
| 89 | except Exception as e: |
| 90 | log.info("Shape inference could not be performed at this time:\n{}".format(e)) |
| 91 | try: |
| 92 | self.graph.fold_constants(fold_shapes=True) |
| 93 | except TypeError as e: |
| 94 | log.error( |