(args)
| 29 | |
| 30 | |
| 31 | def main(args): |
| 32 | # Load saved model |
| 33 | saved_model_path = os.path.realpath(args.saved_model) |
| 34 | assert os.path.isdir(saved_model_path) |
| 35 | graph_def, inputs, outputs = tf_loader.from_saved_model(saved_model_path, None, None, "serve", ["serving_default"]) |
| 36 | with tf.Graph().as_default() as tf_graph: |
| 37 | tf.import_graph_def(graph_def, name="") |
| 38 | with tf_loader.tf_session(graph=tf_graph): |
| 39 | onnx_graph = tfonnx.process_tf_graph(tf_graph, input_names=inputs, output_names=outputs, opset=11) |
| 40 | onnx_model = optimizer.optimize_graph(onnx_graph).make_model("Converted from {}".format(saved_model_path)) |
| 41 | graph = gs.import_onnx(onnx_model) |
| 42 | assert graph |
| 43 | print() |
| 44 | print("ONNX graph created successfully") |
| 45 | |
| 46 | # Set the I/O tensor shapes |
| 47 | graph.inputs[0].shape[0] = args.batch_size |
| 48 | graph.outputs[0].shape[0] = args.batch_size |
| 49 | if args.input_size and args.input_size > 0: |
| 50 | if graph.inputs[0].shape[3] == 3: |
| 51 | # Format NHWC |
| 52 | graph.inputs[0].shape[1] = args.input_size |
| 53 | graph.inputs[0].shape[2] = args.input_size |
| 54 | elif graph.inputs[0].shape[1] == 3: |
| 55 | # Format NCHW |
| 56 | graph.inputs[0].shape[2] = args.input_size |
| 57 | graph.inputs[0].shape[3] = args.input_size |
| 58 | print("ONNX input named '{}' with shape {}".format(graph.inputs[0].name, graph.inputs[0].shape)) |
| 59 | print("ONNX output named '{}' with shape {}".format(graph.outputs[0].name, graph.outputs[0].shape)) |
| 60 | for i in range(4): |
| 61 | if type(graph.inputs[0].shape[i]) != int or graph.inputs[0].shape[i] <= 0: |
| 62 | print("The input shape of the graph is invalid, try overriding it by giving a fixed size with --input_size") |
| 63 | sys.exit(1) |
| 64 | |
| 65 | # Fix Clip Nodes (ReLU6) |
| 66 | for node in [n for n in graph.nodes if n.op == "Clip"]: |
| 67 | for input in node.inputs[1:]: |
| 68 | # In TensorRT, the min/max inputs on a Clip op *must* have fp32 datatype |
| 69 | input.values = np.float32(input.values) |
| 70 | |
| 71 | # Run tensor shape inference |
| 72 | graph.cleanup().toposort() |
| 73 | model = shape_inference.infer_shapes(gs.export_onnx(graph)) |
| 74 | graph = gs.import_onnx(model) |
| 75 | |
| 76 | # Save updated model |
| 77 | graph.cleanup().toposort() |
| 78 | model = gs.export_onnx(graph) |
| 79 | onnx_path = os.path.realpath(args.onnx) |
| 80 | os.makedirs(os.path.dirname(onnx_path), exist_ok=True) |
| 81 | onnx.save(model, onnx_path) |
| 82 | engine_path = os.path.join(os.path.dirname(onnx_path), "engine.trt") |
| 83 | print("ONNX model saved to {}".format(onnx_path)) |
| 84 | |
| 85 | |
| 86 | if __name__ == "__main__": |
no test coverage detected