wehub-resource-sync
101 lines
3.0 KiB
Python
Executable File
101 lines
3.0 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
#
|
|
# 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 onnx
|
|
import torch
|
|
import numpy as np
|
|
import argparse
|
|
import onnx_graphsurgeon as gs
|
|
from post_processing import *
|
|
from packnet_sfm.networks.depth.PackNet01 import PackNet01
|
|
|
|
|
|
def post_process_packnet(model_file, opset=11):
|
|
"""
|
|
Use ONNX graph surgeon to replace upsample and instance normalization nodes. Refer to post_processing.py for details.
|
|
Args:
|
|
model_file : Path to ONNX file
|
|
"""
|
|
# Load the packnet graph
|
|
graph = gs.import_onnx(onnx.load(model_file))
|
|
|
|
if opset >= 11:
|
|
graph = process_pad_nodes(graph)
|
|
|
|
# Replace the subgraph of upsample with a single node with input and scale factor.
|
|
if torch.__version__ < "1.5.0":
|
|
graph = process_upsample_nodes(graph, opset)
|
|
|
|
# Convert the group normalization subgraph into a single plugin node.
|
|
graph = process_groupnorm_nodes(graph)
|
|
|
|
# Remove unused nodes, and topologically sort the graph.
|
|
graph.cleanup().toposort()
|
|
|
|
# Export the onnx graph from graphsurgeon
|
|
onnx.save_model(gs.export_onnx(graph), model_file)
|
|
|
|
print("Saving the ONNX model to {}".format(model_file))
|
|
|
|
|
|
def build_packnet(model_file, args):
|
|
"""
|
|
Construct the packnet network and export it to ONNX
|
|
"""
|
|
input_pyt = torch.randn((1, 3, 192, 640), requires_grad=False)
|
|
|
|
# Build the model
|
|
model_pyt = PackNet01(version="1A")
|
|
|
|
# Convert the model into ONNX
|
|
torch.onnx.export(
|
|
model_pyt, input_pyt, model_file, verbose=args.verbose, opset_version=args.opset
|
|
)
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Exports PackNet01 to ONNX, and post-processes it to insert TensorRT plugins"
|
|
)
|
|
parser.add_argument(
|
|
"-o",
|
|
"--output",
|
|
help="Path to save the generated ONNX model",
|
|
default="model.onnx",
|
|
)
|
|
parser.add_argument(
|
|
"-op", "--opset", type=int, help="ONNX opset to use", default=11
|
|
)
|
|
parser.add_argument(
|
|
"-v",
|
|
"--verbose",
|
|
action="store_true",
|
|
help="Flag to enable verbose logging for torch.onnx.export",
|
|
)
|
|
args = parser.parse_args()
|
|
|
|
# Construct the packnet graph and generate the onnx graph
|
|
build_packnet(args.output, args)
|
|
|
|
# Perform post processing on Instance Normalization and upsampling nodes and create a new ONNX graph
|
|
post_process_packnet(args.output, args.opset)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|