wehub-resource-sync
333 lines
12 KiB
Python
333 lines
12 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 logging
|
|
import numpy as np
|
|
import onnx_graphsurgeon as gs
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logging.getLogger("ModelHelper").setLevel(logging.INFO)
|
|
log = logging.getLogger("ModelHelper")
|
|
|
|
|
|
@gs.Graph.register()
|
|
def op_with_const(self, op, name, input, value):
|
|
"""
|
|
Add an operation with constant to the graph which will operate on the input tensor with the value(s) given.
|
|
:param op: The ONNX operation to perform, i.e. "Add" or "Mul".
|
|
:param input: The tensor to operate on.
|
|
:param value: The value array to operate with.
|
|
:param name: The name to use for the node.
|
|
"""
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created {} node '{}': {}".format(op, name, value.squeeze()))
|
|
const = gs.Constant(name="{}_value:0".format(name), values=value)
|
|
return self.layer(
|
|
name=name, op=op, inputs=[input_tensor, const], outputs=[name + ":0"]
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def matmul(self, name, input, value):
|
|
"""
|
|
Add MatMul operation to the graph which will operate on the input tensor with the value(s) given.
|
|
:param input: The tensor to operate on.
|
|
:param value: The linear transformation matrix to operate with.
|
|
:param name: The name to use for the node.
|
|
"""
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created {} node '{}': {}".format("MatMul", name, value.squeeze()))
|
|
const = gs.Constant(name="{}_value:0".format(name), values=value)
|
|
return self.layer(
|
|
name=name, op="MatMul", inputs=[input_tensor, const], outputs=[name + ":0"]
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def clip(self, name, input, clip_min, clip_max):
|
|
"""
|
|
Add Clip operation to the graph which will operate on the input tensor with the value(s) given.
|
|
:param input: The tensor to operate on.
|
|
:param name: The name to use for the node.
|
|
:param clip_min: Minimum value to include, less is clipped.
|
|
:param clip_max: Maximum value to include, more is clipped.
|
|
"""
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created {} node '{}".format("Clip", name))
|
|
const_min = gs.Constant(
|
|
name="{}_value:0".format(name), values=np.asarray([clip_min], dtype=np.float32)
|
|
)
|
|
const_max = gs.Constant(
|
|
name="{}_value:1".format(name), values=np.asarray([clip_max], dtype=np.float32)
|
|
)
|
|
return self.layer(
|
|
name=name,
|
|
op="Clip",
|
|
inputs=[input_tensor, const_min, const_max],
|
|
outputs=[name + ":0"],
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def slice(self, name, input, starts, ends, axes):
|
|
"""
|
|
Add Slice operation to the graph which will operate on the input tensor with the value(s) given.
|
|
:param op: The ONNX operation to perform, i.e. "Add" or "Mul".
|
|
:param input: The tensor to operate on.
|
|
:param name: The name to use for the node.
|
|
:param starts: Value at which Slice starts.
|
|
:param ends: Value at which Slice ends.
|
|
:param axes: Axes on which Slice operation should be performed.
|
|
"""
|
|
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created {} node '{}".format("Slice", name))
|
|
const_start = gs.Constant(
|
|
name="{}_value:0".format(name), values=np.asarray([starts], dtype=np.int64)
|
|
)
|
|
const_end = gs.Constant(
|
|
name="{}_value:1".format(name), values=np.asarray([ends], dtype=np.int64)
|
|
)
|
|
const_axes = gs.Constant(
|
|
name="{}_value:2".format(name), values=np.asarray([axes], dtype=np.int64)
|
|
)
|
|
return self.layer(
|
|
name=name,
|
|
op="Slice",
|
|
inputs=[input_tensor, const_start, const_end, const_axes],
|
|
outputs=[name + ":0"],
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def unsqueeze(self, name, input, axes=[3]):
|
|
"""
|
|
Adds to the graph an Unsqueeze node for the given axes and to the given input.
|
|
:param self: The gs.Graph object being extended.
|
|
:param name: The name to use for the node.
|
|
:param input: The tensor to be "unsqueezed".
|
|
:param axes: A list of axes on which to add the new dimension(s).
|
|
:return: The first output tensor, to allow chained graph construction.
|
|
"""
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created Unsqueeze node '{}': {}".format(name, axes))
|
|
return self.layer(
|
|
name=name,
|
|
op="Unsqueeze",
|
|
inputs=[input_tensor],
|
|
outputs=[name + ":0"],
|
|
attrs={"axes": axes},
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def squeeze(self, name, input, axes=[2]):
|
|
"""
|
|
Adds to the graph an Squeeze node for the given axes and to the given input.
|
|
:param self: The gs.Graph object being extended.
|
|
:param name: The name to use for the node.
|
|
:param input: The tensor to be "squeezed".
|
|
:param axes: A list of axes on which to remove a dimension(s).
|
|
:return: The first output tensor, to allow chained graph construction.
|
|
"""
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created Squeeze node '{}': {}".format(name, axes))
|
|
return self.layer(
|
|
name=name,
|
|
op="Squeeze",
|
|
inputs=[input_tensor],
|
|
outputs=[name + ":0"],
|
|
attrs={"axes": axes},
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def gather(self, name, data, indices, axes=0):
|
|
"""
|
|
Adds to the graph a Gather node for the given axes and to the given input.
|
|
:param self: The gs.Graph object being extended.
|
|
:param name: The name to use for the node.
|
|
:param data: Data from which to gather specific tensors.
|
|
:param indices: Indices by which to gather data tensors.
|
|
:param axes: A list of axes on which to perform gather operation
|
|
"""
|
|
data_tensor = data if type(data) is gs.Variable else data[0]
|
|
indices_tensor = indices if type(indices) is gs.Variable else indices[0]
|
|
log.debug("Created Gather node '{}': {}".format(name, axes))
|
|
return self.layer(
|
|
name=name,
|
|
op="Gather",
|
|
inputs=[data_tensor, indices_tensor],
|
|
outputs=[name + ":0"],
|
|
attrs={"axes": axes},
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def transpose(self, name, input, perm):
|
|
"""
|
|
Adds to the graph a Transpose node for the given axes permutation and to the given input.
|
|
:param self: The gs.Graph object being extended.
|
|
:param name: The name to use for the node.
|
|
:param input: The tensor to be transposed.
|
|
:param perm: A list of axes defining their order after transposing occurs.
|
|
:return: The first output tensor, to allow chained graph construction.
|
|
"""
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created Transpose node '{}': {}".format(name, perm))
|
|
return self.layer(
|
|
name=name,
|
|
op="Transpose",
|
|
inputs=[input_tensor],
|
|
outputs=[name + ":0"],
|
|
attrs={"perm": perm},
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def sigmoid(self, name, input):
|
|
"""
|
|
Adds to the graph a Sigmoid node for the given input.
|
|
:param self: The gs.Graph object being extended.
|
|
:param name: The name to use for the node.
|
|
:param input: The tensor to be applied to.
|
|
:return: The first output tensor, to allow chained graph construction.
|
|
"""
|
|
input_tensor = input if type(input) is gs.Variable else input[0]
|
|
log.debug("Created Sigmoid node '{}'".format(name))
|
|
return self.layer(
|
|
name=name, op="Sigmoid", inputs=[input_tensor], outputs=[name + ":0"]
|
|
)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def plugin(self, op, name, inputs: list, outputs: list, attrs):
|
|
"""
|
|
Adds to the graph a TensorRT plugin node with the given name, inputs and outputs. The attrs dictionary holds
|
|
attributes to be added to the plugin node.
|
|
:param self: The gs.Graph object being extended.
|
|
:param op: The registered name for the TensorRT plugin.
|
|
:param name: The name to use for the node.
|
|
:param inputs: The list of tensors to use an inputs.
|
|
:param outputs: The list of tensors to use as outputs.
|
|
:param attrs: The dictionary to use as attributes.
|
|
:return: The first output tensor, to allow chained graph construction.
|
|
"""
|
|
log.debug("Created TRT Plugin node '{}': {}".format(name, attrs))
|
|
return self.layer(op=op, name=name, inputs=inputs, outputs=outputs, attrs=attrs)
|
|
|
|
|
|
@gs.Graph.register()
|
|
def find_node_by_op(self, op):
|
|
"""
|
|
Finds the first node in the graph with the given operation name.
|
|
:param self: The gs.Graph object being extended.
|
|
:param op: The operation name to search for.
|
|
:return: The first node matching that performs that op.
|
|
"""
|
|
for node in self.nodes:
|
|
if node.op == op:
|
|
return node
|
|
return None
|
|
|
|
|
|
@gs.Graph.register()
|
|
def find_node_by_op_name(self, op, name):
|
|
"""
|
|
Finds the first node in the graph with the given operation name.
|
|
:param self: The gs.Graph object being extended.
|
|
:param op: The operation name to search for.
|
|
:param name: Selected node name.
|
|
:return: The first node matching that performs that op.
|
|
"""
|
|
for node in self.nodes:
|
|
if node.op == op and node.name == name:
|
|
return node
|
|
return None
|
|
|
|
|
|
@gs.Graph.register()
|
|
def find_node_by_op_input_output_name(
|
|
self, op, input_name, output_name, input_pos=0, output_pos=0
|
|
):
|
|
"""
|
|
Finds the first node in the graph with the given operation name.
|
|
:param self: The gs.Graph object being extended.
|
|
:param op: The operation name to search for.
|
|
:param input_pos: Which input to consider, default is 0.
|
|
:param output_pos: Which output to consider, default is 0.
|
|
:param input_name: Selected input's name.
|
|
:param output_name: Selected output's name.
|
|
:return: The first node matching that performs that op.
|
|
"""
|
|
for node in self.nodes:
|
|
if (
|
|
node.op == op
|
|
and node.inputs[input_pos].name == input_name
|
|
and node.outputs[output_pos].name == output_name
|
|
):
|
|
return node
|
|
return None
|
|
|
|
|
|
@gs.Graph.register()
|
|
def find_descendant_by_op(self, node, op, depth=10):
|
|
"""
|
|
Starting from the given node, finds a node lower in the graph matching the given operation name.
|
|
This is not an exhaustive graph search.
|
|
In order to graph search bfs is used, so runtime complexity is O(V+E).
|
|
:param self: The gs.Graph object being extended.
|
|
:param node: The node to start searching from.
|
|
:param op: The operation name to search for.
|
|
:param depth: Stop searching after traversing these many nodes.
|
|
:return: The first descendant node matching that performs that op.
|
|
"""
|
|
queue = []
|
|
for i in range(depth):
|
|
queue.append(node.o())
|
|
while queue:
|
|
node = queue.pop(0)
|
|
if node.op == op:
|
|
return node
|
|
for child in node.outputs[0].outputs:
|
|
queue.append(child)
|
|
return None
|
|
|
|
|
|
@gs.Graph.register()
|
|
def find_ancestor_by_op(self, node, op, depth=10):
|
|
"""
|
|
Starting from the given node, finds a node higher in the graph matching the given operation name.
|
|
This is not an exhaustive graph search.
|
|
In order to graph search bfs is used, so runtime complexity is O(V+E).
|
|
:param self: The gs.Graph object being extended.
|
|
:param node: The node to start searching from.
|
|
:param op: The operation name to search for.
|
|
:param depth: Stop searching after traversing these many nodes.
|
|
:return: The first ancestor node matching that performs that op.
|
|
"""
|
|
queue = []
|
|
for i in range(depth):
|
|
queue.append(node.i())
|
|
while queue:
|
|
node = queue.pop(0)
|
|
if node.op == op:
|
|
return node
|
|
for child in node.inputs[-1].inputs:
|
|
queue.append(child)
|
|
return None
|