Asfiya Baig
123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
#
|
|
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
#
|
|
|
|
import os
|
|
import sys
|
|
import argparse
|
|
|
|
import onnx
|
|
import onnx_graphsurgeon as gs
|
|
from onnx import shape_inference
|
|
|
|
import numpy as np
|
|
import tensorflow as tf
|
|
from tf2onnx import tfonnx, optimizer, tf_loader
|
|
|
|
|
|
def main(args):
|
|
# Load saved model
|
|
saved_model_path = os.path.realpath(args.saved_model)
|
|
assert os.path.isdir(saved_model_path)
|
|
graph_def, inputs, outputs = tf_loader.from_saved_model(
|
|
saved_model_path, None, None, "serve", ["serving_default"]
|
|
)
|
|
with tf.Graph().as_default() as tf_graph:
|
|
tf.import_graph_def(graph_def, name="")
|
|
with tf_loader.tf_session(graph=tf_graph):
|
|
onnx_graph = tfonnx.process_tf_graph(
|
|
tf_graph, input_names=inputs, output_names=outputs, opset=11
|
|
)
|
|
onnx_model = optimizer.optimize_graph(onnx_graph).make_model(
|
|
"Converted from {}".format(saved_model_path)
|
|
)
|
|
graph = gs.import_onnx(onnx_model)
|
|
assert graph
|
|
print()
|
|
print("ONNX graph created successfully")
|
|
|
|
# Set the I/O tensor shapes
|
|
graph.inputs[0].shape[0] = args.batch_size
|
|
graph.outputs[0].shape[0] = args.batch_size
|
|
if args.input_size and args.input_size > 0:
|
|
if graph.inputs[0].shape[3] == 3:
|
|
# Format NHWC
|
|
graph.inputs[0].shape[1] = args.input_size
|
|
graph.inputs[0].shape[2] = args.input_size
|
|
elif graph.inputs[0].shape[1] == 3:
|
|
# Format NCHW
|
|
graph.inputs[0].shape[2] = args.input_size
|
|
graph.inputs[0].shape[3] = args.input_size
|
|
print(
|
|
"ONNX input named '{}' with shape {}".format(
|
|
graph.inputs[0].name, graph.inputs[0].shape
|
|
)
|
|
)
|
|
print(
|
|
"ONNX output named '{}' with shape {}".format(
|
|
graph.outputs[0].name, graph.outputs[0].shape
|
|
)
|
|
)
|
|
for i in range(4):
|
|
if type(graph.inputs[0].shape[i]) != int or graph.inputs[0].shape[i] <= 0:
|
|
print(
|
|
"The input shape of the graph is invalid, try overriding it by giving a fixed size with --input_size"
|
|
)
|
|
sys.exit(1)
|
|
|
|
# Fix Clip Nodes (ReLU6)
|
|
for node in [n for n in graph.nodes if n.op == "Clip"]:
|
|
for input in node.inputs[1:]:
|
|
# In TensorRT, the min/max inputs on a Clip op *must* have fp32 datatype
|
|
input.values = np.float32(input.values)
|
|
|
|
# Run tensor shape inference
|
|
graph.cleanup().toposort()
|
|
model = shape_inference.infer_shapes(gs.export_onnx(graph))
|
|
graph = gs.import_onnx(model)
|
|
|
|
# Save updated model
|
|
graph.cleanup().toposort()
|
|
model = gs.export_onnx(graph)
|
|
onnx_path = os.path.realpath(args.onnx)
|
|
os.makedirs(os.path.dirname(onnx_path), exist_ok=True)
|
|
onnx.save(model, onnx_path)
|
|
engine_path = os.path.join(os.path.dirname(onnx_path), "engine.trt")
|
|
print("ONNX model saved to {}".format(onnx_path))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument(
|
|
"-m", "--saved_model", help="The TensorFlow saved model directory to load"
|
|
)
|
|
parser.add_argument("-o", "--onnx", help="The output ONNX model file to write")
|
|
parser.add_argument(
|
|
"-b", "--batch_size", type=int, default=1, help="Set the batch size, default: 1"
|
|
)
|
|
parser.add_argument(
|
|
"-i",
|
|
"--input_size",
|
|
type=int,
|
|
help="Override the input height and width, e.g. '380', default: keep original size",
|
|
)
|
|
args = parser.parse_args()
|
|
if not all([args.saved_model, args.onnx]):
|
|
parser.print_help()
|
|
print("\nThese arguments are required: --saved_model and --onnx")
|
|
sys.exit(1)
|
|
main(args)
|