chore: import upstream snapshot with attribution

samples/python/onnx_packnet/post_processing.py

@@ -0,0 +1,248 @@
+#!/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_graphsurgeon as gs
+import argparse
+import onnx
+import numpy as np
+import torch
+
+
+# Pad layer subgraph structure in ONNX (specific to opset 11):
+#               Constant
+#                  |
+#                Shape
+#                  |
+#         Mul   Gather
+#          \     /
+#            Sub
+#             |
+#       ConstantOfShape
+#             |
+#          Concat
+#             |
+#          Reshape
+#             |
+#           Slice
+#             |
+#          Transpose
+#             |
+#          Reshape
+#             |
+#    Input  Cast  Constant
+#       \     |    /
+#            Pad
+def process_pad_nodes(graph):
+    """
+    Fold the pad subgraph into a single layer with pad values as input
+      Input
+       |
+      Pad
+       |
+      Conv
+    """
+    pad_nodes = [node for node in graph.nodes if node.op == "Pad"]
+    for node in pad_nodes:
+        fold_pad_inputs(node, graph)
+
+    return graph
+
+
+def fold_pad_inputs(node, graph):
+    # Gather the amount of padding in each dimension from pytorch graph.
+    if torch.__version__ < "1.5.0":
+        pad_values_pyt = (
+            node.i(1).i(0).i(0).i(0).i(0).i(0).i(0).i(0).attrs["value"].values
+        )
+    elif torch.__version__ < "2.0.0":
+        pad_values_pyt = node.i(1).i(0).i(0).i(0).i(0).i(0).inputs[0].values
+    else:
+        pad_values_pyt = node.i(1).i(0).i(0).i(0).i(0).i(0).i(0).attrs["value"].values
+
+    # Assumption a 4d input tensor
+    onnx_pad_values = [0] * 4 * 2  # 4d tensor and 2 sides padding for each dimension
+    j = 3
+    for i in range(0, len(pad_values_pyt), 2):
+        onnx_pad_values[j] = pad_values_pyt[i]
+        onnx_pad_values[j + 4] = pad_values_pyt[i + 1]
+        j -= 1
+
+    # Change the existing pad tensor to the new onnx_pad values tensor
+    pads_folded_tensor = gs.Constant(
+        name=node.inputs[1].name, values=np.array(onnx_pad_values)
+    )
+    node.inputs[1] = pads_folded_tensor
+
+
+# Pytorch-exported Upsample structure in ONNX:
+#        Mul        Mul
+#         |          |
+#        Cast       Cast
+#         |          |
+#        Floor      Floor
+#         |          |
+#      Unsqueeze  Unsqueeze
+#         \         /
+#           Concat
+#             |
+#            Cast    Cast
+#              \      /
+#                Div
+#                 |
+#     Input     Concat
+#       \         /
+#         Upsample
+def process_upsample_nodes(graph, opset=11):
+    """
+    Replace the upsample structure with structure below
+      Conv   scale_factor
+       |      /
+      Upsample
+       |
+      ReLU
+    """
+    if opset >= 11:
+        upsample_layer_name = "Resize"
+    else:
+        upsample_layer_name = "Upsample"
+
+    upsample_nodes = [node for node in graph.nodes if node.op == upsample_layer_name]
+    for node in upsample_nodes:
+        fold_upsample_inputs(node, graph, opset)
+
+    return graph
+
+
+def fold_upsample_inputs(upsample, graph, opset=11):
+    """
+    Inplace transformation of the graph. The upsample subgraph is collapsed
+    to single upsample node with input and scale factor (constant tensor).
+    Args:
+        upsample: upsample node in the original graph.
+        graph: graph object.
+    """
+
+    if opset == 9:
+        # Gather the scale factor from mul op in the upsample input subgraph
+        scale_factor = (
+            upsample.i(1).i(1).i(0).i(0).i(0).i(0).i(0).i(0).i(1).attrs["value"].values
+        )
+
+        # Create the new scales tensor
+        scales = np.array([1.0, 1.0, scale_factor, scale_factor], dtype=np.float32)
+        scale_tensor = gs.Constant(name=upsample.inputs[-1].name, values=scales)
+
+        # Change the last input to the node to the new constant scales tensor.
+        upsample.inputs[-1] = scale_tensor
+    else:
+        # In opset 11, upsample layer is exported as Resize. We will transform this Resize layer into an Upsample layer
+        # and collapse the input
+        sizes_tensor_name = upsample.inputs[3].name
+
+        # Create the new scales tensor
+        scale_factor = (
+            upsample.i(3).i(1).i().i().i().i().i(0).i(1).attrs["value"].values
+        )
+        scales = np.array([1.0, 1.0, scale_factor, scale_factor], dtype=np.float32)
+        scale_tensor = gs.Constant(name=sizes_tensor_name, values=scales)
+
+        # Rename the Resize op to upsample and add the data and scales as inputs to the upsample layer.
+        input_tensor = upsample.inputs[0]
+        upsample.inputs = [input_tensor, scale_tensor]
+        upsample.op = "Upsample"
+
+
+# Pytorch-exported GroupNorm subgraph in ONNX:
+# Conv
+#   |
+# Reshape    Scale    Bias
+#     \       |       /
+# InstanceNormalization
+#         |
+#      Reshape    Unsqueeze
+#          \      /
+#        Mul (scale)   Unsqueeze
+#           \         /
+#           Add (bias)
+#              |
+#            ReLU
+def process_groupnorm_nodes(graph):
+    """
+    Gather the instance normalization nodes and the rest of the subgraph
+    and convert into a single group normalization node.
+    """
+    instancenorms = [node for node in graph.nodes if node.op == "InstanceNormalization"]
+    for node in instancenorms:
+        convert_to_groupnorm(node, graph)
+
+    return graph
+
+
+def retrieve_attrs(instancenorm):
+    """
+    Gather the required attributes for the GroupNorm plugin from the subgraph.
+    Args:
+        instancenorm: Instance Normalization node in the graph.
+    """
+    attrs = {}
+    # The 2nd dimension of the Reshape shape is the number of groups
+    attrs["num_groups"] = instancenorm.i().i(1).attrs["value"].values[1]
+    attrs["eps"] = instancenorm.attrs["epsilon"]
+
+    # 1 is the default plugin version the parser will search for, and therefore can be omitted,
+    # but we include it here for illustrative purposes.
+    attrs["plugin_version"] = "1"
+
+    # "" is the default plugin namespace the parser will use, included here for illustrative purposes
+    attrs["plugin_namespace"] = ""
+
+    return attrs
+
+
+def convert_to_groupnorm(instancenorm, graph):
+    """
+    Convert the Pytorch-exported GroupNorm subgraph to the subgraph below
+    Conv
+      |
+    GroupNorm
+      |
+    ReLU
+    Attributes:
+        instancenorm: Instance Normalization node in the graph.
+        graph: Input graph object
+    """
+    # Retrieve the instancenorm attributes and create the replacement node
+    attrs = retrieve_attrs(instancenorm)
+    groupnorm = gs.Node(op="GroupNormalizationPlugin", attrs=attrs)
+    graph.nodes.append(groupnorm)
+
+    # The plugin needs to receive an input from the Conv node, and output to the ReLU node
+    conv_output_tensor = instancenorm.i().inputs[0]  # Output of Conv
+    relu_input_tensor = instancenorm.o().o().o().outputs[0]  # Output of Add
+
+    # Reconnect inputs/outputs to the groupnorm plugin
+    conv_output_tensor.outputs[0] = groupnorm
+    relu_input_tensor.inputs[0] = groupnorm
+
+    # Add scale and bias constant tensors to group norm plugin
+    if torch.__version__ < "1.5.0":
+        groupnorm.inputs.append(instancenorm.o().o().i(1).inputs[0])
+        groupnorm.inputs.append(instancenorm.o().o().o().i(1).inputs[0])
+    else:
+        groupnorm.inputs.append(instancenorm.o().o().inputs[1])
+        groupnorm.inputs.append(instancenorm.o().o().o().inputs[1])